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/Cargo.lock b/Cargo.lock index 43862759c38..c58c85dfd99 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9752,6 +9752,7 @@ dependencies = [ "codspeed-divan-compat", "insta", "itertools 0.14.0", + "num-traits", "pco", "rand 0.10.2", "rstest", 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..11a7134d5da --- /dev/null +++ b/vortex-array/src/patches_v2.rs @@ -0,0 +1,753 @@ +// 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::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use num_traits::AsPrimitive; +use num_traits::PrimInt; +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; + +/// How patch values combine with already-decoded base values. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum PatchMode { + /// The patch value replaces the base value. + #[default] + Replace, + /// The patch value is bitwise-ORed into the base value, supplying low bits the base does not + /// store (FloatQuant-style: the base holds the value with its low bits shifted away and + /// shifted back up on decode, so most elements need no patch). + CombineLow, + /// The patch value is shifted left by `shift` and bitwise-ORed into the base value, supplying + /// high bits above a `shift`-bit base (block-residual-style: the base stores only the low + /// `shift` bits and outliers patch their high bits back in). + CombineHigh { + /// Number of low bits held by the base value. + shift: u8, + }, +} + +static REPLACE_APPLIES: AtomicU64 = AtomicU64::new(0); +static COMBINE_LOW_APPLIES: AtomicU64 = AtomicU64::new(0); +static COMBINE_HIGH_APPLIES: AtomicU64 = AtomicU64::new(0); + +/// Returns how many [`PatchesV2::apply_into`] calls ran per mode: +/// `(replace, combine_low, combine_high)`. +pub fn patches_v2_mode_applies() -> (u64, u64, u64) { + ( + REPLACE_APPLIES.load(Ordering::Relaxed), + COMBINE_LOW_APPLIES.load(Ordering::Relaxed), + COMBINE_HIGH_APPLIES.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, + /// How patch values combine with base values on decode. + mode: PatchMode, +} + +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, + mode: PatchMode::Replace, + } + } + + /// Returns this patch set with the given combine mode. + pub fn with_mode(mut self, mode: PatchMode) -> Self { + self.mode = mode; + self + } + + /// Returns how patch values combine with base values on decode. + pub fn mode(&self) -> PatchMode { + self.mode + } + + /// 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`]. + /// + /// Only [`PatchMode::Replace`] patch sets convert: [`Patches`] has no combine semantics. + pub fn to_patches(&self, ctx: &mut ExecutionCtx) -> VortexResult { + vortex_ensure!( + self.mode == PatchMode::Replace, + "only Replace-mode PatchesV2 can convert to Patches, got {:?}", + self.mode + ); + 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 into `out` according to this patch set's [`PatchMode`]. + /// + /// `out` holds the decoded base values; each patched position is overwritten (Replace) or + /// bitwise-OR-combined with the (optionally shifted) patch value (CombineLow / CombineHigh). + pub fn apply_into(&self, out: &mut [T], ctx: &mut ExecutionCtx) -> VortexResult<()> + where + T: crate::dtype::NativePType + PrimInt, + { + vortex_ensure!( + out.len() == self.array_len, + "PatchesV2 apply_into expects {} elements, got {}", + self.array_len, + out.len() + ); + let values = self.values.clone().execute::(ctx)?; + let values = values.as_slice::(); + match self.mode { + PatchMode::Replace => { + REPLACE_APPLIES.fetch_add(1, Ordering::Relaxed); + self.apply_each(ctx, |logical, ordinal| out[logical] = values[ordinal]) + } + PatchMode::CombineLow => { + COMBINE_LOW_APPLIES.fetch_add(1, Ordering::Relaxed); + self.apply_each(ctx, |logical, ordinal| { + out[logical] = out[logical] | values[ordinal] + }) + } + PatchMode::CombineHigh { shift } => { + COMBINE_HIGH_APPLIES.fetch_add(1, Ordering::Relaxed); + self.apply_each(ctx, |logical, ordinal| { + out[logical] = out[logical] | (values[ordinal] << usize::from(shift)) + }) + } + } + } + + /// 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(), + ) + } + .with_mode(self.mode), + )) + } + + /// 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 apply_into_modes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Patches at rows 1 and 3 of a 4-row array, patch values 0b01 and 0b10. + let indices = PrimitiveArray::new(buffer![1u64, 3], Validity::NonNullable); + let values = PrimitiveArray::new(buffer![0b01u64, 0b10], 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 = [0b100u64; 4]; + patches.apply_into(&mut out, &mut ctx)?; + assert_eq!(out, [0b100, 0b01, 0b100, 0b10]); + + let mut out = [0b100u64; 4]; + patches + .clone() + .with_mode(PatchMode::CombineLow) + .apply_into(&mut out, &mut ctx)?; + assert_eq!(out, [0b100, 0b101, 0b100, 0b110]); + + let mut out = [0b100u64; 4]; + patches + .clone() + .with_mode(PatchMode::CombineHigh { shift: 3 }) + .apply_into(&mut out, &mut ctx)?; + assert_eq!(out, [0b100, 0b1100, 0b100, 0b10100]); + + // Combine-mode patch sets no longer round-trip into global Patches. + assert!( + patches + .with_mode(PatchMode::CombineLow) + .to_patches(&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(()) + } +} diff --git a/vortex-bench/src/datasets/taxi_data.rs b/vortex-bench/src/datasets/taxi_data.rs index 919a0fa4dff..98e96a9ff4b 100644 --- a/vortex-bench/src/datasets/taxi_data.rs +++ b/vortex-bench/src/datasets/taxi_data.rs @@ -122,3 +122,110 @@ pub async fn taxi_data_vortex_compact() -> Result { }) .await } + +#[cfg(test)] +mod tests { + use vortex::array::ArrayEq; + use vortex::array::Canonical; + use vortex::array::EqMode; + use vortex::array::VortexSessionExecute; + use vortex::array::arrays::Struct; + use vortex::array::arrays::chunked::ChunkedArrayExt; + use vortex::array::arrays::struct_::StructArrayExt; + use vortex::array::patches_v2::patches_v2_mode_applies; + use vortex::compressor::BtrBlocksCompressorBuilder; + use vortex::compressor::FloatQuantV2Scheme; + use vortex::compressor::SchemeExt; + use vortex::compressor::float_quant_v2_encode_counts; + + use super::*; + use crate::conversions::parquet_to_vortex_chunks; + + /// Compress the taxi dataset with the default schemes (FloatQuantV2 enabled), decode it back, + /// and report which patch combine modes were exercised. + /// + /// Every column is also compressed by a control compressor with FloatQuantV2 excluded, so a + /// column that fails equality under both compressors points at the comparison or another + /// scheme, not at FloatQuantV2. + /// + /// Ignored: downloads the taxi parquet file. Run with: + /// `cargo test -p vortex-bench --release taxi_float_quant -- --ignored --nocapture` + #[tokio::test(flavor = "multi_thread")] + #[ignore] + async fn taxi_float_quant_v2_report() -> Result<()> { + let chunks = parquet_to_vortex_chunks(taxi_data_parquet().await?).await?; + let compressor = BtrBlocksCompressorBuilder::default().build(); + let control_compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([FloatQuantV2Scheme.id()]) + .build(); + let mut ctx = SESSION.create_execution_ctx(); + + let mut plain_bytes = 0u64; + let mut compressed_bytes = 0u64; + let mut control_bytes = 0u64; + let mut quant_nodes = 0usize; + for idx in 0..chunks.nchunks() { + let chunk = chunks.chunk(idx); + let compressed = compressor.compress(chunk, &mut ctx)?; + let control_compressed = control_compressor.compress(chunk, &mut ctx)?; + plain_bytes += chunk.nbytes() as u64; + compressed_bytes += compressed.nbytes() as u64; + control_bytes += control_compressed.nbytes() as u64; + quant_nodes += compressed + .tree_display() + .to_string() + .matches("float_quant_v2") + .count(); + + let expected = chunk.clone().execute::(&mut ctx)?.into_array(); + let decoded = compressed.execute::(&mut ctx)?.into_array(); + let control_decoded = control_compressed + .execute::(&mut ctx)? + .into_array(); + let expected = expected.as_::(); + let decoded = decoded.as_::(); + let control_decoded = control_decoded.as_::(); + for (field_idx, name) in expected.names().iter().enumerate() { + // Struct canonicalization is shallow: execute each field so encoded children + // (including FloatQuantV2) actually decode. + let exp = expected + .unmasked_field(field_idx) + .clone() + .execute::(&mut ctx)? + .into_array(); + let quant = decoded + .unmasked_field(field_idx) + .clone() + .execute::(&mut ctx)? + .into_array(); + let control = control_decoded + .unmasked_field(field_idx) + .clone() + .execute::(&mut ctx)? + .into_array(); + let quant_eq = exp.array_eq(&quant, EqMode::Value); + let control_eq = exp.array_eq(&control, EqMode::Value); + if quant_eq != control_eq { + println!( + "chunk {idx} field {name}: quant_eq={quant_eq} control_eq={control_eq}" + ); + } + assert!( + quant_eq || !control_eq, + "chunk {idx} field {name} mismatches only with FloatQuantV2 enabled" + ); + } + } + + let (low_encodes, high_encodes) = float_quant_v2_encode_counts(); + let (replace_applies, low_applies, high_applies) = patches_v2_mode_applies(); + println!("taxi: plain={plain_bytes} quant={compressed_bytes} control={control_bytes}"); + println!("float_quant_v2 nodes selected: {quant_nodes}"); + println!("float_quant_v2 encodes: combine_low={low_encodes} combine_high={high_encodes}"); + println!( + "patches_v2 applies: replace={replace_applies} combine_low={low_applies} \ + combine_high={high_applies}" + ); + Ok(()) + } +} diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 4e22f042adf..2773ee110e6 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -18,6 +18,7 @@ all-features = true [dependencies] itertools = { workspace = true } +num-traits = { workspace = true } pco = { workspace = true, optional = true } rand = { workspace = true } vortex-alp = { workspace = true } @@ -33,6 +34,7 @@ vortex-onpair = { workspace = true, optional = true } vortex-pco = { workspace = true, optional = true } vortex-runend = { workspace = true } vortex-sequence = { workspace = true } +vortex-session = { workspace = true } vortex-sparse = { workspace = true } vortex-utils = { workspace = true } vortex-zigzag = { workspace = true } diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 6f38e29cd86..f29b319a5aa 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -44,6 +44,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ //////////////////////////////////////////////////////////////////////////////////////////////// &float::ALPScheme, &float::ALPRDScheme, + &float::FloatQuantV2Scheme, &float::FloatDictScheme, &float::NullDominatedSparseScheme, &float::FloatRLEScheme, 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/float/mod.rs b/vortex-btrblocks/src/schemes/float/mod.rs index 1301184ac0c..4f12d6f7467 100644 --- a/vortex-btrblocks/src/schemes/float/mod.rs +++ b/vortex-btrblocks/src/schemes/float/mod.rs @@ -5,6 +5,7 @@ mod alp; mod alprd; +mod quant_v2; mod rle; mod sparse; @@ -15,6 +16,13 @@ pub use alp::ALPScheme; pub use alprd::ALPRDScheme; #[cfg(feature = "pco")] pub use pco::PcoScheme; +pub use quant_v2::FloatQuantV2; +pub use quant_v2::FloatQuantV2Array; +pub use quant_v2::FloatQuantV2Scheme; +pub use quant_v2::FloatQuantV2Slots; +pub use quant_v2::QuantMode; +pub use quant_v2::float_quant_v2_encode; +pub use quant_v2::float_quant_v2_encode_counts; pub use rle::FloatRLEScheme; pub use sparse::NullDominatedSparseScheme; // Re-export builtin schemes from vortex-compressor. diff --git a/vortex-btrblocks/src/schemes/float/quant_v2.rs b/vortex-btrblocks/src/schemes/float/quant_v2.rs new file mode 100644 index 00000000000..96467e61074 --- /dev/null +++ b/vortex-btrblocks/src/schemes/float/quant_v2.rs @@ -0,0 +1,685 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Float quantization with sparse combine-mode [`PatchesV2`] exceptions. +//! +//! `FloatQuantV2` splits each float's raw bits at a chosen bit position `k` and stores only one +//! side densely; the other side is zero for most elements and the exceptions are patched back in +//! with a combine-mode [`PatchesV2`] set: +//! +//! - [`PatchMode::CombineLow`]: the dense `primary` child holds `bits >> k` (an integer that +//! compresses well); elements whose low `k` bits are non-zero store them as patch values that +//! are bitwise-ORed in after the primary is shifted back up. This wins on decimal-style floats +//! whose mantissas end in zero bits (the FloatQuant split from the pco-adjacent encodings PR). +//! - [`PatchMode::CombineHigh`]: the dense `primary` child holds `bits & ((1 << k) - 1)` +//! (bitpackable to `k` bits); elements with any bits above `k` store `bits >> k` as patch +//! values ORed back in shifted left by `k`. This wins on columns dominated by `0.0` with sparse +//! outliers (the block-residual high-bit split). +//! +//! This is a prototype encoding co-located with its compression scheme; it has no slice or +//! serde-registry integration yet and exists to measure the combine-mode patch layouts on real +//! data. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use num_traits::PrimInt as NumPrimInt; +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::Canonical; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::array_slots; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::patches::PATCH_CHUNK_SIZE; +use vortex_array::patches_v2::PatchMode; +use vortex_array::patches_v2::PatchesV2; +use vortex_array::scalar::Scalar; +use vortex_array::serde::ArrayChildren; +use vortex_array::smallvec::smallvec; +use vortex_array::validity::Validity; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::ValidityChild; +use vortex_array::vtable::ValidityVTableFromChild; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; + +static LOW_ENCODES: AtomicU64 = AtomicU64::new(0); +static HIGH_ENCODES: AtomicU64 = AtomicU64::new(0); + +/// Returns how many arrays were encoded per split: `(combine_low, combine_high)`. +pub fn float_quant_v2_encode_counts() -> (u64, u64) { + ( + LOW_ENCODES.load(Ordering::Relaxed), + HIGH_ENCODES.load(Ordering::Relaxed), + ) +} + +/// A [`FloatQuantV2`]-encoded Vortex array. +pub type FloatQuantV2Array = Array; + +/// Marker type for the FloatQuantV2 encoding. +#[derive(Clone, Debug)] +pub struct FloatQuantV2; + +/// Child slots of a [`FloatQuantV2Array`]. +#[array_slots(FloatQuantV2)] +pub struct FloatQuantV2Slots { + /// The dense side of the split: `bits >> k` (low mode) or `bits & ((1 << k) - 1)` (high + /// mode), as unsigned integers carrying the float's validity. + #[slot(0)] + pub primary: ArrayRef, + /// Chunk-local `u16` patch indices. + #[slot(1)] + pub patch_indices: ArrayRef, + /// Patch values: the missing low or high bits per exception. + #[slot(2)] + pub patch_values: ArrayRef, + /// `u32` prefix patch counts per chunk with a leading zero. + #[slot(3)] + pub patch_chunk_offsets: ArrayRef, +} + +/// Which side of the bit split the sparse patches carry. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum QuantMode { + /// Primary stores `bits >> k`; patches OR the low `k` bits back in. + Low, + /// Primary stores the low `k` bits; patches OR `bits >> k` back in shifted by `k`. + High, +} + +#[derive(Clone, Debug)] +pub struct FloatQuantV2Data { + mode: QuantMode, + k: u8, +} + +impl Display for FloatQuantV2Data { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let mode = match self.mode { + QuantMode::Low => "low", + QuantMode::High => "high", + }; + write!(f, "mode: {mode}, k: {}", self.k) + } +} + +impl ArrayHash for FloatQuantV2Data { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.mode.hash(state); + self.k.hash(state); + } +} + +impl ArrayEq for FloatQuantV2Data { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.mode == other.mode && self.k == other.k + } +} + +impl FloatQuantV2 { + /// Construct a new array from the split components. + pub fn try_new( + primary: ArrayRef, + patch_indices: ArrayRef, + patch_values: ArrayRef, + patch_chunk_offsets: ArrayRef, + mode: QuantMode, + k: u8, + ) -> VortexResult { + let dtype = logical_dtype(primary.dtype())?; + let len = primary.len(); + vortex_ensure!( + patch_indices.len() == patch_values.len(), + "FloatQuantV2 patch indices and values must have the same length" + ); + vortex_ensure!(!patch_indices.is_empty(), "FloatQuantV2 requires patches"); + let data = FloatQuantV2Data { mode, k }; + let slots = smallvec![ + Some(primary), + Some(patch_indices), + Some(patch_values), + Some(patch_chunk_offsets), + ]; + Ok(unsafe { + Array::from_parts_unchecked( + ArrayParts::new(FloatQuantV2, dtype, len, data).with_slots(slots), + ) + }) + } +} + +fn logical_dtype(primary_dtype: &DType) -> VortexResult { + match primary_dtype { + DType::Primitive(PType::U32, n) => Ok(DType::Primitive(PType::F32, *n)), + DType::Primitive(PType::U64, n) => Ok(DType::Primitive(PType::F64, *n)), + d => vortex_bail!(MismatchedTypes: "u32 or u64", d), + } +} + +/// Rebuild the combine-mode patch set from the array's children. +fn patches(array: ArrayView<'_, FloatQuantV2>) -> PatchesV2 { + let mode = match array.mode { + QuantMode::Low => PatchMode::CombineLow, + QuantMode::High => PatchMode::CombineHigh { shift: array.k }, + }; + // SAFETY: The encoder builds valid chunk-local components and the array is immutable. + unsafe { + PatchesV2::new_unchecked( + array.as_ref().len(), + 0, + array.patch_indices().clone(), + array.patch_values().clone(), + array.patch_chunk_offsets().clone(), + ) + } + .with_mode(mode) +} + +impl VTable for FloatQuantV2 { + type TypedArrayData = FloatQuantV2Data; + + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.float_quant_v2"); + *ID + } + + fn validate( + &self, + data: &FloatQuantV2Data, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let view = FloatQuantV2SlotsView::from_slots(slots); + let expected = logical_dtype(view.primary.dtype())?; + vortex_ensure!(dtype == &expected, "expected dtype {expected}, got {dtype}"); + vortex_ensure!( + view.primary.len() == len, + "expected len {len}, got {}", + view.primary.len() + ); + vortex_ensure!( + view.patch_indices.dtype() == &DType::Primitive(PType::U16, NonNullable), + "FloatQuantV2 patch indices must be non-nullable u16" + ); + vortex_ensure!( + view.patch_chunk_offsets.len() == len.div_ceil(PATCH_CHUNK_SIZE) + 1, + "FloatQuantV2 chunk offsets have the wrong length" + ); + let width = u8::try_from(view.primary.dtype().as_ptype().bit_width())?; + vortex_ensure!( + data.k > 0 && data.k < width, + "FloatQuantV2 split position {} out of range for {width}-bit floats", + data.k + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("FloatQuantV2 buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + let mode = match array.mode { + QuantMode::Low => 0u8, + QuantMode::High => 1u8, + }; + Ok(Some(vec![mode, array.k])) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + _buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + metadata.len() == 2, + "FloatQuantV2 expects 2 metadata bytes, got {}", + metadata.len() + ); + let mode = match metadata[0] { + 0 => QuantMode::Low, + 1 => QuantMode::High, + other => vortex_bail!("invalid FloatQuantV2 mode {other}"), + }; + let k = metadata[1]; + let primary_dtype = match dtype { + DType::Primitive(PType::F32, n) => DType::Primitive(PType::U32, *n), + DType::Primitive(PType::F64, n) => DType::Primitive(PType::U64, *n), + d => vortex_bail!(MismatchedTypes: "f32 or f64", d), + }; + let value_dtype = DType::Primitive(primary_dtype.as_ptype(), NonNullable); + let primary = children.get(0, &primary_dtype, len)?; + let patch_len = children + .get(1, &DType::Primitive(PType::U16, NonNullable), 0)? + .len(); + let patch_indices = + children.get(1, &DType::Primitive(PType::U16, NonNullable), patch_len)?; + let patch_values = children.get(2, &value_dtype, patch_len)?; + let patch_chunk_offsets = children.get( + 3, + &DType::Primitive(PType::U32, NonNullable), + len.div_ceil(PATCH_CHUNK_SIZE) + 1, + )?; + let slots = smallvec![ + Some(primary), + Some(patch_indices), + Some(patch_values), + Some(patch_chunk_offsets), + ]; + Ok(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + FloatQuantV2Data { mode, k }, + ) + .with_slots(slots)) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + FloatQuantV2Slots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let primary = array.primary().clone().execute::(ctx)?; + let validity = primary.validity()?; + let patches = patches(array.as_view()); + let k = usize::from(array.k); + let decoded = match primary.ptype() { + PType::U32 => decode::(&primary, &patches, array.mode, k, validity, ctx)?, + PType::U64 => decode::(&primary, &patches, array.mode, k, validity, ctx)?, + p => vortex_bail!("invalid FloatQuantV2 primary ptype {p}"), + }; + Ok(ExecutionResult::done(decoded.into_array())) + } +} + +fn decode( + primary: &PrimitiveArray, + patches: &PatchesV2, + mode: QuantMode, + k: usize, + validity: Validity, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut bits = BufferMut::::with_capacity(primary.len()); + match mode { + QuantMode::Low => bits.extend(primary.as_slice::().iter().map(|&p| p << k)), + QuantMode::High => bits.extend_from_slice(primary.as_slice::()), + } + patches.apply_into(bits.as_mut_slice(), ctx)?; + let float_ptype = match U::PTYPE { + PType::U32 => PType::F32, + PType::U64 => PType::F64, + p => vortex_bail!("invalid FloatQuantV2 primary ptype {p}"), + }; + Ok(PrimitiveArray::new(bits.freeze(), validity).reinterpret_cast(float_ptype)) +} + +impl OperationsVTable for FloatQuantV2 { + fn scalar_at( + array: ArrayView<'_, FloatQuantV2>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let primary = array.primary().execute_scalar(index, ctx)?; + if primary.is_null() { + return Ok(Scalar::null(array.dtype().clone())); + } + let k = u32::from(array.k); + let mut bits = primary + .as_primitive() + .as_::() + .vortex_expect("non-null primary"); + if array.mode == QuantMode::Low { + bits <<= k; + } + let patch_set = patches(array); + if let Some(ordinal) = patch_set.search_index(index, ctx)?.to_found() { + let patch = array + .patch_values() + .execute_scalar(ordinal, ctx)? + .as_primitive() + .as_::() + .vortex_expect("non-null patch value"); + bits |= match array.mode { + QuantMode::Low => patch, + QuantMode::High => patch << k, + }; + } + Ok(match array.dtype().as_ptype() { + PType::F32 => Scalar::primitive( + f32::from_bits(u32::try_from(bits)?), + array.dtype().nullability(), + ), + PType::F64 => Scalar::primitive(f64::from_bits(bits), array.dtype().nullability()), + p => vortex_bail!("invalid FloatQuantV2 ptype {p}"), + }) + } +} + +impl ValidityChild for FloatQuantV2 { + fn validity_child(array: ArrayView<'_, FloatQuantV2>) -> ArrayRef { + array.primary().clone() + } +} + +/// One candidate split: mode, split position, and estimated encoded size in bits. +#[derive(Clone, Copy)] +struct Split { + mode: QuantMode, + k: u8, + exceptions: usize, + cost_bits: u128, +} + +/// Encode a float array, returning `None` when no bit split beats storing the raw bits. +pub fn float_quant_v2_encode( + array: ArrayView<'_, vortex_array::arrays::Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let _ = ctx; + match array.ptype() { + PType::F32 => encode_typed::(array, 32), + PType::F64 => encode_typed::(array, 64), + _ => Ok(None), + } +} + +fn encode_typed( + array: ArrayView<'_, vortex_array::arrays::Primitive>, + width: u32, +) -> VortexResult> { + let bits_array = array.reinterpret_cast(U::PTYPE); + let bits = bits_array.as_slice::(); + let len = bits.len(); + if len < PATCH_CHUNK_SIZE { + return Ok(None); + } + + // Histogram trailing zeros and bit lengths in one pass. `tz` of zero is `width`. + let word_bits = width as usize; + let mut tz_hist = vec![0usize; word_bits + 1]; + let mut bl_hist = vec![0usize; word_bits + 1]; + for &b in bits { + let tz = b.trailing_zeros().min(width) as usize; + let bl = (width - b.leading_zeros().min(width)) as usize; + tz_hist[tz] += 1; + bl_hist[bl] += 1; + } + // Suffix sums: elements with at least k trailing zeros / at most k significant bits. + let mut tz_at_least = vec![0usize; word_bits + 2]; + for k in (0..=word_bits).rev() { + tz_at_least[k] = tz_at_least[k + 1] + tz_hist[k]; + } + let mut bl_at_most = vec![0usize; word_bits + 1]; + let mut acc = 0usize; + for k in 0..=word_bits { + acc += bl_hist[k]; + bl_at_most[k] = acc; + } + + let chunk_bits = (len.div_ceil(PATCH_CHUNK_SIZE) + 1) as u128 * 32; + let patch_bits = 16 + width as u128; + let max_exceptions = len / 16; + let mut best: Option = None; + for k in 1..word_bits { + for (mode, exceptions, dense_bits) in [ + ( + QuantMode::Low, + len - tz_at_least[k], + (word_bits - k) as u128, + ), + (QuantMode::High, len - bl_at_most[k], k as u128), + ] { + if exceptions == 0 || exceptions > max_exceptions { + continue; + } + let cost_bits = len as u128 * dense_bits + exceptions as u128 * patch_bits + chunk_bits; + if best.is_none_or(|b| cost_bits < b.cost_bits) { + best = Some(Split { + mode, + k: u8::try_from(k)?, + exceptions, + cost_bits, + }); + } + } + } + let Some(split) = best else { + return Ok(None); + }; + // Require a real win over the raw bits before taking on the split. + if split.cost_bits * 16 >= len as u128 * width as u128 * 15 { + return Ok(None); + } + + let k = usize::from(split.k); + let low_mask = (U::one() << k) - U::one(); + let mut primary = BufferMut::::with_capacity(len); + let mut locals = BufferMut::::with_capacity(split.exceptions); + let mut values = BufferMut::::with_capacity(split.exceptions); + let mut chunk_offsets = vec![0u32; len.div_ceil(PATCH_CHUNK_SIZE) + 1]; + for (i, &b) in bits.iter().enumerate() { + let (dense, patch) = match split.mode { + QuantMode::Low => (b >> k, b & low_mask), + QuantMode::High => (b & low_mask, b >> k), + }; + primary.push(dense); + if patch != U::zero() { + locals.push(u16::try_from(i % PATCH_CHUNK_SIZE)?); + values.push(patch); + chunk_offsets[i / PATCH_CHUNK_SIZE + 1] += 1; + } + } + for chunk in 0..chunk_offsets.len() - 1 { + chunk_offsets[chunk + 1] += chunk_offsets[chunk]; + } + if locals.is_empty() { + return Ok(None); + } + match split.mode { + QuantMode::Low => LOW_ENCODES.fetch_add(1, Ordering::Relaxed), + QuantMode::High => HIGH_ENCODES.fetch_add(1, Ordering::Relaxed), + }; + + let validity = array.validity()?; + Some(FloatQuantV2::try_new( + PrimitiveArray::new(primary.freeze(), validity).into_array(), + PrimitiveArray::new(locals.freeze(), Validity::NonNullable).into_array(), + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + PrimitiveArray::new(Buffer::from(chunk_offsets), Validity::NonNullable).into_array(), + split.mode, + split.k, + )) + .transpose() +} + +/// Float quantization with sparse combine-mode patches. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct FloatQuantV2Scheme; + +impl Scheme for FloatQuantV2Scheme { + fn scheme_name(&self) -> &'static str { + "vortex.float.quant_v2" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches!( + canonical.dtype(), + DType::Primitive(PType::F32 | PType::F64, _) + ) + } + + fn produced_encodings(&self) -> Vec { + vec![FloatQuantV2.id()] + } + + /// Children: primary=0. + fn num_children(&self) -> usize { + 1 + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + // The dense primary is the same size as the input until integer compression shrinks it. + if compress_ctx.finished_cascading() { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + CompressionEstimate::Deferred(DeferredEstimate::Sample) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let Some(encoded) = float_quant_v2_encode(data.array_as_primitive(), exec_ctx)? else { + // No viable split: hand back the input so sampling never selects this scheme. + return Ok(data.array_as_primitive().array().clone()); + }; + let compressed_primary = compressor.compress_child( + encoded.primary(), + &compress_ctx, + crate::SchemeExt::id(self), + 0, + exec_ctx, + )?; + Ok(FloatQuantV2::try_new( + compressed_primary, + encoded.patch_indices().clone(), + encoded.patch_values().clone(), + encoded.patch_chunk_offsets().clone(), + encoded.mode, + encoded.k, + )? + .into_array()) + } +} + +#[cfg(test)] +mod tests { + use std::f64::consts::PI; + + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::assert_arrays_eq; + + use super::*; + + #[test] + fn low_split_roundtrip() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Decimal-style doubles with zero low mantissa bits, plus sparse full-precision values. + let values: Vec = (0..4096) + .map(|i| { + if i % 700 == 3 { + PI + i as f64 + } else { + (i / 4) as f64 + 0.5 + } + }) + .collect(); + let array = PrimitiveArray::from_iter(values); + let encoded = + float_quant_v2_encode(array.as_view(), &mut ctx)?.vortex_expect("split found"); + assert_eq!(encoded.mode, QuantMode::Low); + let decoded = encoded.into_array().execute::(&mut ctx)?; + assert_arrays_eq!(array.into_array(), decoded.into_array(), &mut ctx); + Ok(()) + } + + #[test] + fn high_split_roundtrip() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Mostly zeros with sparse arbitrary doubles: the high bits are patched back in. + let values: Vec = (0..4096) + .map(|i| if i % 900 == 7 { 1.75 + i as f64 } else { 0.0 }) + .collect(); + let array = PrimitiveArray::from_iter(values); + let encoded = + float_quant_v2_encode(array.as_view(), &mut ctx)?.vortex_expect("split found"); + assert_eq!(encoded.mode, QuantMode::High); + let decoded = encoded + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_arrays_eq!(array.clone().into_array(), decoded.into_array(), &mut ctx); + + // scalar_at agrees with the decoded values. + for index in [0, 7, 907, 4095] { + assert_eq!( + encoded.as_view().array().execute_scalar(index, &mut ctx)?, + Scalar::primitive(array.as_slice::()[index], NonNullable) + ); + } + Ok(()) + } +} 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(()) +} diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 1d1ab1252ac..7e31140653d 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -146,7 +146,10 @@ pub mod compressor { pub use vortex_btrblocks::BtrBlocksCompressor; pub use vortex_btrblocks::BtrBlocksCompressorBuilder; pub use vortex_btrblocks::Scheme; + pub use vortex_btrblocks::SchemeExt; pub use vortex_btrblocks::SchemeId; + pub use vortex_btrblocks::schemes::float::FloatQuantV2Scheme; + pub use vortex_btrblocks::schemes::float::float_quant_v2_encode_counts; } /// Vortex editions: named, frozen sets of encodings with a read-compatibility guarantee.