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(()) +} 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(()) +} 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()) +}