diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 9085390b67b..74108e54134 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -70,3 +70,7 @@ harness = false name = "cast_bitpacked" harness = false required-features = ["_test-harness"] + +[[bench]] +name = "chunked_decompress" +harness = false diff --git a/encodings/fastlanes/benches/chunked_decompress.rs b/encodings/fastlanes/benches/chunked_decompress.rs new file mode 100644 index 00000000000..2a6867b4c4f --- /dev/null +++ b/encodings/fastlanes/benches/chunked_decompress.rs @@ -0,0 +1,358 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measures the composable chunked-decompression vtable path (`decompress_chunks`) for +//! FoR-over-BitPacked against: +//! +//! - `two_pass`: full fused decompression to a `PrimitiveArray` followed by a second pass over +//! the materialized buffer (the pattern `decompress_chunks` is designed to replace). +//! - `hand_fused`: a monomorphized loop over `BitUnpackedChunks` with the reference folded in +//! algebraically — the no-dispatch upper bound, equivalent to the fused `unpack_map` kernels. +//! +//! The gap between `chunked_vtable` and `hand_fused` is the exact price of the generic +//! mechanism: one virtual call per 1024-element chunk per encoding level plus one in-place +//! add pass over the L1-resident chunk at the FoR level. + +// The `#[gat(Item)]` import expansion trips unused_imports even though the lending iterator +// machinery requires it. +#![allow(unused_imports)] + +use std::hint::black_box; +use std::ops::Range; +use std::sync::LazyLock; + +use divan::Bencher; +use lending_iterator::gat; +use lending_iterator::prelude::Item; +#[gat(Item)] +use lending_iterator::prelude::LendingIterator; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::chunk_iter::ChunkMut; +use vortex_array::chunk_iter::ChunkSink; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::FoR; +use vortex_fastlanes::FoRArraySlotsExt; +use vortex_fastlanes::bitpack_compress::bitpack_encode; +use vortex_fastlanes::unpack_iter::BitUnpackedChunks; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); + +const LEN: usize = 4 * 1024 * 1024; +const REFERENCE: u32 = 1_000_000; + +/// FoR(reference=1M) over BitPacked(bit_width=10), no patches. +fn make_for_bitpacked() -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + let deltas = PrimitiveArray::from_iter( + (0..LEN as u64).map(|i| u32::try_from((i * 7) % 1000).vortex_expect("fits")), + ); + let bp = bitpack_encode(&deltas, 10, None, &mut ctx).vortex_expect("bench"); + FoR::try_new(bp.into_array(), Scalar::from(REFERENCE)) + .vortex_expect("bench") + .into_array() +} + +struct SumSink { + total: u64, +} + +impl ChunkSink for SumSink { + #[inline] + fn accept(&mut self, chunk: ChunkMut<'_>, _row_range: Range) -> VortexResult<()> { + self.total = self + .total + .wrapping_add(chunk.as_slice::().iter().map(|&v| v as u64).sum()); + Ok(()) + } +} + +/// Sum through the composable vtable path. With an unsigned reference over BitPacked, FoR takes +/// its fused streaming path: each block is unpacked with the reference folded into the kernel +/// and handed straight to the sink — one streaming pass, no materialization, no extra add pass. +#[divan::bench] +fn chunked_vtable_sum(bencher: Bencher) { + let array = make_for_bitpacked(); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + let mut sink = SumSink { total: 0 }; + array + .decompress_chunks(&mut ctx, &mut sink) + .vortex_expect("bench"); + black_box(sink.total) + }); +} + +/// Same as `chunked_vtable_sum` but with a signed FoR reference, which skips the fused +/// FoR+BitPacked dispatch and exercises the generic composition: BitPacked streams plain +/// unpacked blocks and FoR's sink adapter adds the reference in place per chunk. +#[divan::bench] +fn chunked_vtable_sum_generic_compose(bencher: Bencher) { + let mut ctx = SESSION.create_execution_ctx(); + let deltas = PrimitiveArray::from_iter( + (0..LEN as u64).map(|i| i32::try_from((i * 7) % 1000).vortex_expect("fits")), + ); + let bp = bitpack_encode(&deltas, 10, None, &mut ctx).vortex_expect("bench"); + let array = FoR::try_new(bp.into_array(), Scalar::from(-1_000_000i32)) + .vortex_expect("bench") + .into_array(); + + struct SumSinkI32 { + total: i64, + } + impl ChunkSink for SumSinkI32 { + #[inline] + fn accept(&mut self, chunk: ChunkMut<'_>, _row_range: Range) -> VortexResult<()> { + self.total = self + .total + .wrapping_add(chunk.as_slice::().iter().map(|&v| v as i64).sum()); + Ok(()) + } + } + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + let mut sink = SumSinkI32 { total: 0 }; + array + .decompress_chunks(&mut ctx, &mut sink) + .vortex_expect("bench"); + black_box(sink.total) + }); +} + +/// Sum by fully decompressing (the fused FoR+BitPacked execute path) and then re-reading the +/// materialized array: two passes over `LEN` values. +#[divan::bench] +fn two_pass_sum(bencher: Bencher) { + let array = make_for_bitpacked(); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + let primitive = array + .execute::(&mut ctx) + .vortex_expect("bench"); + let total: u64 = primitive + .as_slice::() + .iter() + .map(|&v| v as u64) + .fold(0, u64::wrapping_add); + black_box(total) + }); +} + +/// Monomorphized upper bound: iterate the BitPacked child's unpacked blocks directly and fold +/// the FoR reference in algebraically — no dynamic dispatch, no extra in-place pass. +#[divan::bench] +fn hand_fused_sum(bencher: Bencher) { + let array = make_for_bitpacked(); + bencher.with_inputs(|| array.clone()).bench_values(|array| { + let for_ = array.as_::(); + let bp = for_.encoded().as_::(); + let mut chunks: BitUnpackedChunks = bp.unpacked_chunks().vortex_expect("bench"); + let mut total = 0u64; + let mut count = 0u64; + if let Some(initial) = chunks.initial() { + total = total.wrapping_add(initial.iter().map(|&v| v as u64).sum()); + count += initial.len() as u64; + } + let mut iter = chunks.full_chunks(); + while let Some(chunk) = iter.next() { + total = total.wrapping_add(chunk.iter().map(|&v| v as u64).sum()); + count += chunk.len() as u64; + } + if let Some(trailer) = chunks.trailer() { + total = total.wrapping_add(trailer.iter().map(|&v| v as u64).sum()); + count += trailer.len() as u64; + } + black_box(total.wrapping_add(count * REFERENCE as u64)) + }); +} + +/// Monomorphized version of exactly what the vtable path does per chunk (unpack, in-place +/// reference add, then fold) but with zero dynamic dispatch. The gap between this and +/// `hand_fused_sum` is the cost of the extra in-place pass; the gap between this and +/// `chunked_vtable_sum` is the pure dynamic-dispatch cost (two virtual calls per 1024 values). +#[divan::bench] +fn hand_chunked_add_pass_sum(bencher: Bencher) { + let array = make_for_bitpacked(); + bencher.with_inputs(|| array.clone()).bench_values(|array| { + let for_ = array.as_::(); + let bp = for_.encoded().as_::(); + let mut chunks: BitUnpackedChunks = bp.unpacked_chunks().vortex_expect("bench"); + let mut total = 0u64; + let mut fold = |chunk: &mut [u32]| { + for v in chunk.iter_mut() { + *v = v.wrapping_add(REFERENCE); + } + total = total.wrapping_add(chunk.iter().map(|&v| v as u64).sum()); + }; + if let Some(initial) = chunks.initial() { + fold(initial); + } + let mut iter = chunks.full_chunks(); + while let Some(chunk) = iter.next() { + fold(chunk); + } + if let Some(trailer) = chunks.trailer() { + fold(trailer); + } + black_box(total) + }); +} + +struct WriteSink<'a> { + out: &'a mut Vec, +} + +impl ChunkSink for WriteSink<'_> { + #[inline] + fn accept(&mut self, chunk: ChunkMut<'_>, _row_range: Range) -> VortexResult<()> { + self.out.extend_from_slice(chunk.as_slice::()); + Ok(()) + } +} + +/// Materialize through the chunked path (unpack to scratch, add reference in place, copy out) — +/// upper-bounds the cost of the non-fused composition when the consumer wants a full buffer. +#[divan::bench] +fn chunked_vtable_decompress_into(bencher: Bencher) { + let array = make_for_bitpacked(); + bencher + .with_inputs(|| { + ( + array.clone(), + SESSION.create_execution_ctx(), + Vec::::with_capacity(LEN), + ) + }) + .bench_values(|(array, mut ctx, mut out)| { + let mut sink = WriteSink { out: &mut out }; + array + .decompress_chunks(&mut ctx, &mut sink) + .vortex_expect("bench"); + black_box(out.len()) + }); +} + +/// Materialize through the fused execute path (`decode_into` writes unpacked+shifted values +/// straight into the output buffer) — the specialized baseline for full decompression. +#[divan::bench] +fn fused_decompress(bencher: Bencher) { + let array = make_for_bitpacked(); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + let primitive = array + .execute::(&mut ctx) + .vortex_expect("bench"); + black_box(primitive.len()) + }); +} + +/// Patched(Constant): the base never materializes — Constant re-emits one L1-resident scratch +/// chunk and Patched's sink adapter overwrites the sparse patch rows per chunk. +fn make_patched_constant() -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + let inner = vortex_array::arrays::ConstantArray::new(42u32, LEN).into_array(); + let n_patches = LEN / 1000; + let patches = vortex_array::patches::Patches::new( + LEN, + 0, + PrimitiveArray::from_iter((0..n_patches as u64).map(|i| i * 1000)).into_array(), + PrimitiveArray::from_iter( + (0..n_patches as u64).map(|i| u32::try_from(i % 90_000).vortex_expect("fits")), + ) + .into_array(), + None, + ) + .vortex_expect("bench"); + vortex_array::arrays::Patched::from_array_and_patches(inner, &patches, &mut ctx) + .vortex_expect("bench") + .into_array() +} + +#[divan::bench] +fn chunked_vtable_sum_patched_constant(bencher: Bencher) { + let array = make_patched_constant(); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + let mut sink = SumSink { total: 0 }; + array + .decompress_chunks(&mut ctx, &mut sink) + .vortex_expect("bench"); + black_box(sink.total) + }); +} + +#[divan::bench] +fn two_pass_sum_patched_constant(bencher: Bencher) { + let array = make_patched_constant(); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + let primitive = array + .execute::(&mut ctx) + .vortex_expect("bench"); + let total: u64 = primitive + .as_slice::() + .iter() + .map(|&v| v as u64) + .fold(0, u64::wrapping_add); + black_box(total) + }); +} + +/// Sparse decompression baseline: `execute` on Patched(Constant) canonicalizes the constant into +/// a full-length buffer, then scatters the sparse patches over it — one full-buffer write pass +/// plus sparse writes, materializing the array. +#[divan::bench] +fn sparse_decompress_patched_constant(bencher: Bencher) { + let array = make_patched_constant(); + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| { + let primitive = array + .execute::(&mut ctx) + .vortex_expect("bench"); + black_box(primitive.len()) + }); +} + +/// Materialize Patched(Constant) through the chunked path: splat the constant into an 8KB +/// scratch, patch it, and copy each chunk out to the destination buffer. +#[divan::bench] +fn chunked_decompress_into_patched_constant(bencher: Bencher) { + let array = make_patched_constant(); + bencher + .with_inputs(|| { + ( + array.clone(), + SESSION.create_execution_ctx(), + Vec::::with_capacity(LEN), + ) + }) + .bench_values(|(array, mut ctx, mut out)| { + let mut sink = WriteSink { out: &mut out }; + array + .decompress_chunks(&mut ctx, &mut sink) + .vortex_expect("bench"); + black_box(out.len()) + }); +} diff --git a/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs b/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs new file mode 100644 index 00000000000..a5ae674ebcc --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs @@ -0,0 +1,234 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Streaming chunked decompression for bit-packed arrays. +//! +//! This backs [`VTable::decompress_chunks`](vortex_array::vtable::VTable::decompress_chunks) for +//! [`BitPacked`]: each FastLanes block is unpacked into the decompressor's cache-resident scratch +//! buffer, patches falling in that block are applied in place, and the block is handed to the +//! sink — the array is never materialized in full. + +use num_traits::AsPrimitive; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::chunk_iter::ChunkMut; +use vortex_array::chunk_iter::ChunkSink; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PhysicalPType; +use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::patches::Patches; +use vortex_error::VortexResult; + +use crate::BitPacked; +use crate::BitPackedArrayExt; +use crate::unpack_iter::BitPacked as BitPackedUnpack; +use crate::unpack_iter::UnpackStrategy; +use crate::unpack_iter::UnpackedChunks; + +pub(crate) fn decompress_chunks( + array: ArrayView<'_, BitPacked>, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, +) -> VortexResult<()> { + match_each_integer_ptype!(array.as_ref().dtype().as_ptype(), |T| { + decompress_chunks_typed::(array, ctx, sink) + }) +} + +fn decompress_chunks_typed( + array: ArrayView<'_, BitPacked>, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, +) -> VortexResult<()> { + if array.as_ref().is_empty() { + return Ok(()); + } + + let patch_list = match array.patches() { + None => Vec::new(), + Some(patches) => build_patch_list(&patches, ctx, |v: T| v)?, + }; + + let mut chunks = array.unpacked_chunks::()?; + stream_unpacked_chunks(&mut chunks, &patch_list, sink) +} + +/// Materialize sparse patches once as sorted (local row, value) pairs so the per-chunk loop only +/// advances a cursor, applying `map` to each patch value. This is the only heap state the +/// streaming path allocates, and it is proportional to the patch count, not the array length. +pub(crate) fn build_patch_list( + patches: &Patches, + ctx: &mut ExecutionCtx, + map: impl Fn(T) -> T, +) -> VortexResult> { + let indices = patches.indices().clone().execute::(ctx)?; + let values = patches.values().clone().execute::(ctx)?; + let values = values.as_slice::(); + let offset = patches.offset(); + Ok(match_each_unsigned_integer_ptype!(indices.ptype(), |P| { + indices + .as_slice::

() + .iter() + .zip(values) + .map(|(&idx, &v)| (

>::as_(idx) - offset, map(v))) + .collect() + })) +} + +/// Walk every unpacked FastLanes block in order, patch it in place from the pre-built cursor +/// list, and hand it to the sink. Generic over the [`UnpackStrategy`] so fused strategies (e.g. +/// FoR's reference-add unpack) stream through the same loop with zero extra passes. +pub(crate) fn stream_unpacked_chunks>( + chunks: &mut UnpackedChunks, + patch_list: &[(usize, T)], + sink: &mut dyn ChunkSink, +) -> VortexResult<()> { + let mut patch_cursor = 0usize; + let mut result = Ok(()); + chunks.for_each_unpacked_chunk(|chunk, range| { + if result.is_err() { + return; + } + while let Some(&(row, value)) = patch_list.get(patch_cursor) + && row < range.end + { + chunk[row - range.start] = value; + patch_cursor += 1; + } + result = sink.accept(ChunkMut::new(chunk), range); + }); + result +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::chunk_iter::ChunkMut; + use vortex_array::dtype::NativePType; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::BitPackedArrayExt; + use crate::FoRArraySlotsExt; + use crate::FoRData; + use crate::bitpack_compress::bitpack_encode; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + fn collect_chunks(array: &vortex_array::ArrayRef) -> VortexResult> { + let mut ctx = SESSION.create_execution_ctx(); + let mut out = Vec::with_capacity(array.len()); + array.decompress_chunks_or_materialize(&mut ctx, &mut |chunk: ChunkMut<'_>, + _range: std::ops::Range< + usize, + >| + -> VortexResult<()> { + out.extend_from_slice(chunk.as_slice::()); + Ok(()) + })?; + Ok(out) + } + + fn assert_chunks_match_execute( + array: vortex_array::ArrayRef, + ) -> VortexResult<()> { + let chunked = collect_chunks::(&array)?; + let mut ctx = SESSION.create_execution_ctx(); + let expected = array.execute::(&mut ctx)?; + assert_eq!(chunked.as_slice(), expected.as_slice::()); + Ok(()) + } + + #[test] + fn bitpacked_chunks_match_execute() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter((0..5000u32).map(|i| i % 900)); + let bp = bitpack_encode(&values, 10, None, &mut ctx)?; + let array = bp.into_array(); + assert!(array.supports_decompress_chunks()); + assert_chunks_match_execute::(array) + } + + #[test] + fn bitpacked_chunks_with_patches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter( + (0..5000u32).map(|i| if i % 700 == 0 { 100_000 + i } else { i % 900 }), + ); + let bp = bitpack_encode(&values, 10, None, &mut ctx)?; + assert!(bp.patches().is_some()); + assert_chunks_match_execute::(bp.into_array()) + } + + #[test] + fn bitpacked_chunks_sliced() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter( + (0..5000u32).map(|i| if i % 700 == 0 { 100_000 + i } else { i % 900 }), + ); + let bp = bitpack_encode(&values, 10, None, &mut ctx)?.into_array(); + // Slice crossing chunk boundaries with a non-zero offset. + let sliced = bp.slice(517..4013)?; + // The slice may no longer be a BitPacked array head, but chunked iteration must still + // stream correct values via whatever encodings the slice resolves to. + assert_chunks_match_execute::(sliced) + } + + #[test] + fn for_over_bitpacked_chunks() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter((0..5000i64).map(|i| 1_000_000 + (i * 7) % 800)); + let for_array = FoRData::encode(values, &mut ctx)?; + assert!( + for_array.encoded().as_opt::().is_some() + || for_array + .encoded() + .as_opt::() + .is_some() + ); + assert_chunks_match_execute::(for_array.into_array()) + } + + /// An unsigned reference over a BitPacked child takes the fused `FoRStrategy` streaming + /// path (reference folded into the unpack kernel), including patch handling. + #[test] + fn for_over_bitpacked_fused_chunks_with_patches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let deltas = PrimitiveArray::from_iter( + (0..5000u32).map(|i| if i % 700 == 0 { 100_000 + i } else { i % 900 }), + ); + let bp = bitpack_encode(&deltas, 10, None, &mut ctx)?; + assert!(bp.patches().is_some()); + let for_array = crate::FoR::try_new( + bp.into_array(), + vortex_array::scalar::Scalar::from(1_000_000u32), + )?; + assert_chunks_match_execute::(for_array.into_array()) + } + + #[test] + fn fallback_primitive_chunks() -> VortexResult<()> { + let values = PrimitiveArray::from_iter(0..3000i32).into_array(); + assert_chunks_match_execute::(values) + } + + #[test] + fn empty_array_chunks() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter(Vec::::new()); + let bp = bitpack_encode(&values, 0, None, &mut ctx)?; + let chunked = collect_chunks::(&bp.into_array())?; + assert!(chunked.is_empty()); + Ok(()) + } +} diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 03fa3ed7f4c..e84a3900e5a 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -26,6 +26,7 @@ use vortex_error::vortex_err; pub mod bitpack_compress; pub mod bitpack_decompress; +pub(crate) mod chunked_decompress; pub mod unpack_iter; use crate::BitPackedArray; diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index efa0677a91e..b2cd687c9ac 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -9,6 +9,7 @@ pub use array::BitPackedDataParts; pub use array::BitPackedSlots; pub use array::bitpack_compress; pub use array::bitpack_decompress; +pub(crate) use array::chunked_decompress; pub use array::unpack_iter; pub(crate) mod compute; diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 68fbf1b41d3..995d68b7e0d 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -19,6 +19,7 @@ use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; +use vortex_array::chunk_iter::ChunkSink; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; @@ -49,6 +50,7 @@ use crate::bitpack_decompress::unpack_into_primitive_builder; use crate::bitpacking::array::BitPackedSlots; use crate::bitpacking::array::BitPackedSlotsView; use crate::bitpacking::array::PATCH_SLOTS; +use crate::bitpacking::chunked_decompress; use crate::bitpacking::vtable::rules::RULES; mod kernels; mod operations; @@ -297,6 +299,18 @@ impl VTable for BitPacked { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } + + fn supports_decompress_chunks(_array: ArrayView<'_, Self>) -> bool { + true + } + + fn decompress_chunks( + array: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, + ) -> VortexResult<()> { + chunked_decompress::decompress_chunks(array, ctx, sink) + } } #[derive(Clone, Debug)] diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index cb94618c071..1f68b11bc6b 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; + use fastlanes::FoR; use num_traits::PrimInt; use num_traits::WrappingAdd; @@ -8,6 +10,8 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::arrays::PrimitiveArray; use vortex_array::builders::PrimitiveBuilder; +use vortex_array::chunk_iter::ChunkMut; +use vortex_array::chunk_iter::ChunkSink; use vortex_array::dtype::NativePType; use vortex_array::dtype::PhysicalPType; use vortex_array::dtype::UnsignedPType; @@ -21,6 +25,7 @@ use crate::BitPacked; use crate::BitPackedArrayExt; use crate::FoRArray; use crate::bitpack_decompress; +use crate::bitpacking::chunked_decompress; use crate::r#for::array::FoRArrayExt; use crate::r#for::array::FoRArraySlotsExt; use crate::unpack_iter::UnpackStrategy; @@ -137,6 +142,107 @@ pub(crate) fn fused_decompress< Ok(builder.finish_into_primitive()) } +/// Whether [`decompress_chunks`] can stream: either the fused FoR+BitPacked path applies, or the +/// encoded child itself supports streaming (generic composition). +pub(crate) fn supports_decompress_chunks(array: ArrayView<'_, crate::FoR>) -> bool { + (array.reference_scalar().dtype().is_unsigned_int() && array.encoded().is::()) + || array.encoded().supports_decompress_chunks() +} + +/// Streaming chunked decompression for FoR arrays. +/// +/// When the child is a [`BitPacked`] array (the common layout), this streams through the fused +/// [`FoRStrategy`] unpack kernel — the reference is folded into `unchecked_unfor_pack` itself, +/// exactly like [`fused_decompress`], so there is no extra add pass at all. +/// +/// Otherwise FoR composes generically over any encoded child: each chunk the child streams up is +/// shifted by the reference value in place (one pass over an L1-resident chunk) and forwarded to +/// the downstream sink. The adapter lives on this stack frame — no heap state is added on the +/// way down. +pub(crate) fn decompress_chunks( + array: ArrayView<'_, crate::FoR>, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, +) -> VortexResult<()> { + // Fused fast path, mirroring the dispatch in `decompress`. + if array.reference_scalar().dtype().is_unsigned_int() + && let Some(bp) = array.encoded().as_opt::() + { + return match_each_unsigned_integer_ptype!(array.ptype(), |T| { + fused_decompress_chunks::(array, bp, ctx, sink) + }); + } + + match_each_integer_ptype!(array.ptype(), |T| { + let reference = array + .reference_scalar() + .as_primitive() + .typed_value::() + .vortex_expect("reference must be non-null"); + if reference == 0 { + array.encoded().decompress_chunks(ctx, sink) + } else { + let mut adapter = AddReferenceSink { + reference, + inner: sink, + }; + array.encoded().decompress_chunks(ctx, &mut adapter) + } + }) +} + +/// Stream FoR-over-BitPacked chunks through the fused unpack kernel: each FastLanes block is +/// unpacked with the reference added by `unchecked_unfor_pack`, patched in place (patch values +/// get the reference applied up front), and handed to the sink. +fn fused_decompress_chunks + UnsignedPType + FoR + WrappingAdd>( + for_: ArrayView<'_, crate::FoR>, + bp: ArrayView<'_, BitPacked>, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, +) -> VortexResult<()> { + if bp.as_ref().is_empty() { + return Ok(()); + } + + let ref_ = for_ + .reference_scalar() + .as_primitive() + .as_::() + .vortex_expect("cannot be null"); + + let mut unpacked = UnpackedChunks::try_new_with_strategy( + FoRStrategy { reference: ref_ }, + bp.packed().as_host().clone(), + bp.bit_width() as usize, + bp.offset() as usize, + bp.len(), + )?; + + let patch_list = match bp.patches() { + None => Vec::new(), + Some(patches) => { + chunked_decompress::build_patch_list(&patches, ctx, |v: T| v.wrapping_add(&ref_))? + } + }; + + chunked_decompress::stream_unpacked_chunks(&mut unpacked, &patch_list, sink) +} + +struct AddReferenceSink<'a, T> { + reference: T, + inner: &'a mut dyn ChunkSink, +} + +impl ChunkSink for AddReferenceSink<'_, T> { + #[inline] + fn accept(&mut self, mut chunk: ChunkMut<'_>, row_range: Range) -> VortexResult<()> { + for v in chunk.as_slice_mut::() { + *v = v.wrapping_add(&self.reference); + } + self.inner.accept(chunk, row_range) + } +} + fn decompress_primitive( values: Buffer, min: T, diff --git a/encodings/fastlanes/src/for/vtable/mod.rs b/encodings/fastlanes/src/for/vtable/mod.rs index 452abcd4562..0d8187dce5b 100644 --- a/encodings/fastlanes/src/for/vtable/mod.rs +++ b/encodings/fastlanes/src/for/vtable/mod.rs @@ -18,6 +18,7 @@ use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; +use vortex_array::chunk_iter::ChunkSink; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -36,6 +37,7 @@ use crate::FoRData; use crate::r#for::array::FoRArrayExt; use crate::r#for::array::FoRSlots; use crate::r#for::array::FoRSlotsView; +use crate::r#for::array::for_decompress; use crate::r#for::array::for_decompress::decompress; use crate::r#for::vtable::rules::PARENT_RULES; @@ -161,6 +163,18 @@ impl VTable for FoR { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(decompress(&array, ctx)?.into_array())) } + + fn supports_decompress_chunks(array: ArrayView<'_, Self>) -> bool { + for_decompress::supports_decompress_chunks(array) + } + + fn decompress_chunks( + array: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, + ) -> VortexResult<()> { + for_decompress::decompress_chunks(array, ctx, sink) + } } #[derive(Clone, Debug)] diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 8b9b2812bfa..9a7d021ead7 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -220,6 +220,17 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { ctx: &mut ExecutionCtx, ) -> VortexResult; + /// Returns whether this encoding (recursively) supports streaming chunked decompression. + fn supports_decompress_chunks(&self, this: &ArrayRef) -> bool; + + /// Stream the array's decompressed values through `sink` in cache-resident chunks. + fn decompress_chunks( + &self, + this: &ArrayRef, + ctx: &mut ExecutionCtx, + sink: &mut dyn crate::chunk_iter::ChunkSink, + ) -> VortexResult<()>; + /// Execute the scalar at the given index. /// /// This method panics if the index is out of bounds for the array. @@ -490,6 +501,21 @@ impl DynArrayData for ArrayData { V::execute(typed, ctx) } + fn supports_decompress_chunks(&self, this: &ArrayRef) -> bool { + let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; + V::supports_decompress_chunks(view) + } + + fn decompress_chunks( + &self, + this: &ArrayRef, + ctx: &mut ExecutionCtx, + sink: &mut dyn crate::chunk_iter::ChunkSink, + ) -> VortexResult<()> { + let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; + V::decompress_chunks(view, ctx, sink) + } + fn execute_scalar( &self, this: &ArrayRef, diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 3078a6ae9f5..8252f1b1aa7 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -187,6 +187,46 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { /// incorrectly contains null values. fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult; + /// Returns whether this encoding can stream decompressed chunks via + /// [`decompress_chunks`](Self::decompress_chunks) without materializing anything. + /// + /// Defaults to `false`: streaming is an explicit capability, never a silent fallback. Leaf + /// encodings that override `decompress_chunks` return `true`; wrapper encodings must + /// propagate the check into the children their implementation streams from (e.g. via + /// [`ArrayRef::supports_decompress_chunks`](crate::ArrayRef::supports_decompress_chunks)), + /// so support of the whole tree is decided before any chunk is emitted. + fn supports_decompress_chunks(array: ArrayView<'_, Self>) -> bool { + _ = array; + false + } + + /// Stream the array's decompressed values through `sink` in cache-resident chunks. + /// + /// See [`chunk_iter`](crate::chunk_iter) for the contract and cost model. The whole point of + /// this method is to decompress and transform block-by-block while the block is L1-resident; + /// implementations must therefore stream directly out of their decompression kernel (leaf + /// encodings) or interpose a stack-allocated [`ChunkSink`](crate::chunk_iter::ChunkSink) + /// adapter and recurse into their child via + /// [`ArrayRef::decompress_chunks`](crate::ArrayRef::decompress_chunks) (wrappers) — never + /// materialize the array. Encodings that cannot do this leave the default, which reports + /// unsupported; callers wanting a materializing fallback opt in by name via + /// [`ArrayRef::decompress_chunks_or_materialize`](crate::ArrayRef::decompress_chunks_or_materialize). + /// + /// Only called when [`supports_decompress_chunks`](Self::supports_decompress_chunks) returned + /// `true`; the caller guarantees the array is primitive-typed. Implementations must emit + /// contiguous, in-order chunks covering exactly `0..len` (checked in debug builds). + fn decompress_chunks( + array: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + sink: &mut dyn crate::chunk_iter::ChunkSink, + ) -> VortexResult<()> { + _ = (ctx, sink); + vortex_bail!( + "decompress_chunks is not supported by encoding {}", + array.encoding_id() + ) + } + /// Attempt to reduce the array to a simpler representation without changing logical values. /// /// Reductions are opportunistic and may return `Ok(None)` when no cheaper representation is diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 6917d979239..7b97e3a9478 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -42,6 +42,9 @@ use crate::builders::PrimitiveBuilder; use crate::builders::VarBinViewBuilder; use crate::builders::builder_with_capacity; use crate::canonical::Canonical; +use crate::chunk_iter::ChunkMut; +use crate::chunk_iter::ChunkSink; +use crate::chunk_iter::DECOMPRESS_CHUNK_LEN; use crate::dtype::DType; use crate::dtype::OffsetBuilderPType; use crate::match_each_decimal_value; @@ -186,6 +189,38 @@ impl VTable for Constant { )?)) } + fn supports_decompress_chunks(_array: ArrayView<'_, Self>) -> bool { + true + } + + fn decompress_chunks( + array: ArrayView<'_, Self>, + _ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, + ) -> VortexResult<()> { + let len = array.as_ref().len(); + // The entry point guarantees a primitive dtype. A null constant streams unspecified (but + // initialized) values, matching the not-streamed-validity contract. + match_each_native_ptype!(array.as_ref().dtype().as_ptype(), |T| { + let value: T = array + .scalar() + .as_primitive() + .typed_value::() + .unwrap_or_default(); + let mut scratch = [value; DECOMPRESS_CHUNK_LEN]; + let mut start = 0; + while start < len { + let n = (len - start).min(DECOMPRESS_CHUNK_LEN); + // Sinks may mutate the chunk in place (e.g. Patched patching over it), so the + // scratch is refilled before every emission. + scratch[..n].fill(value); + sink.accept(ChunkMut::new(&mut scratch[..n]), start..start + n)?; + start += n; + } + Ok(()) + }) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index e2473c19749..a7a45e4424f 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -45,6 +45,8 @@ use crate::arrays::primitive::PrimitiveDataParts; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; +use crate::chunk_iter::ChunkMut; +use crate::chunk_iter::ChunkSink; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::PType; @@ -309,6 +311,111 @@ impl VTable for Patched { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } + + fn supports_decompress_chunks(array: ArrayView<'_, Self>) -> bool { + array.inner().supports_decompress_chunks() + } + + fn decompress_chunks( + array: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, + ) -> VortexResult<()> { + if array.as_ref().is_empty() { + return Ok(()); + } + + let len = array.as_ref().len(); + let offset = array.offset(); + let n_lanes = array.n_lanes(); + + // The patch children are tiny relative to the array; materialize them once up front and + // flatten the lane-transposed layout into row-sorted (row, value) pairs so the per-chunk + // wrapper below only advances a cursor. + let lane_offsets = array + .lane_offsets() + .clone() + .execute::(ctx)?; + let indices = array + .patch_indices() + .clone() + .execute::(ctx)?; + let values = array + .patch_values() + .clone() + .execute::(ctx)?; + + match_each_native_ptype!(values.ptype(), |V| { + let mut adapter = PatchChunkSink { + patch_list: build_row_sorted_patches::( + lane_offsets.as_slice::(), + indices.as_slice::(), + values.as_slice::(), + offset, + len, + n_lanes, + ), + cursor: 0, + inner: sink, + }; + array.inner().decompress_chunks(ctx, &mut adapter) + }) + } +} + +/// Flatten the lane-transposed patch layout into row-sorted (row, value) pairs so the streaming +/// chunk wrapper only advances a cursor. +fn build_row_sorted_patches( + lane_offsets: &[u32], + indices: &[u16], + values: &[V], + offset: usize, + len: usize, + n_lanes: usize, +) -> Vec<(usize, V)> { + let mut patch_list: Vec<(usize, V)> = Vec::with_capacity(values.len()); + let n_chunks = (offset + len).div_ceil(1024); + for chunk in 0..n_chunks { + let start = lane_offsets[chunk * n_lanes] as usize; + let stop = lane_offsets[chunk * n_lanes + n_lanes] as usize; + for idx in start..stop { + // The indices slice is measured as an offset into the 1024-value chunk. + let index = chunk * 1024 + indices[idx] as usize; + if index < offset || index >= offset + len { + continue; + } + patch_list.push((index - offset, values[idx])); + } + } + // Patches are sorted by (chunk, lane), not by row; the stable sort preserves the original + // application order for any duplicate rows. + patch_list.sort_by_key(|&(row, _)| row); + patch_list +} + +/// Sink adapter that overwrites patched rows in each streamed chunk before forwarding it. +struct PatchChunkSink<'a, V> { + patch_list: Vec<(usize, V)>, + cursor: usize, + inner: &'a mut dyn ChunkSink, +} + +impl ChunkSink for PatchChunkSink<'_, V> { + #[inline] + fn accept( + &mut self, + mut chunk: ChunkMut<'_>, + row_range: std::ops::Range, + ) -> VortexResult<()> { + let out = chunk.as_slice_mut::(); + while let Some(&(row, value)) = self.patch_list.get(self.cursor) + && row < row_range.end + { + out[row - row_range.start] = value; + self.cursor += 1; + } + self.inner.accept(chunk, row_range) + } } /// Apply patches on top of the existing value types. diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 45256a7a498..08d070edb29 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -17,6 +17,9 @@ use crate::arrays::primitive::PrimitiveData; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; +use crate::chunk_iter::ChunkMut; +use crate::chunk_iter::ChunkSink; +use crate::chunk_iter::DECOMPRESS_CHUNK_LEN; use crate::dtype::DType; use crate::dtype::PType; use crate::match_each_native_ptype; @@ -181,6 +184,30 @@ impl VTable for Primitive { Ok(ExecutionResult::done(array)) } + fn supports_decompress_chunks(_array: ArrayView<'_, Self>) -> bool { + true + } + + fn decompress_chunks( + array: ArrayView<'_, Self>, + _ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, + ) -> VortexResult<()> { + // Already decompressed: stream the buffer through one reusable L1-resident scratch chunk + // (chunks are handed out mutably, so the shared buffer cannot be exposed directly). + match_each_native_ptype!(array.ptype(), |P| { + let values = array.as_slice::

(); + let mut scratch = vec![P::default(); values.len().min(DECOMPRESS_CHUNK_LEN)]; + for (i, chunk) in values.chunks(DECOMPRESS_CHUNK_LEN).enumerate() { + let start = i * DECOMPRESS_CHUNK_LEN; + let scratch = &mut scratch[..chunk.len()]; + scratch.copy_from_slice(chunk); + sink.accept(ChunkMut::new(scratch), start..start + chunk.len())?; + } + Ok(()) + }) + } + fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, diff --git a/vortex-array/src/chunk_iter.rs b/vortex-array/src/chunk_iter.rs new file mode 100644 index 00000000000..735115b5fa0 --- /dev/null +++ b/vortex-array/src/chunk_iter.rs @@ -0,0 +1,375 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Streaming chunked decompression. +//! +//! [`ArrayRef::decompress_chunks`] walks an array's decompressed values in cache-resident chunks +//! (~1024 elements) without materializing the full array. Each encoding can override +//! [`VTable::decompress_chunks`](crate::vtable::VTable::decompress_chunks) to stream chunks +//! straight out of its decompression kernel; wrapper encodings (e.g. FoR) compose by interposing +//! a stack-allocated [`ChunkSink`] adapter around the downstream sink and recursing into their +//! child through the erased [`ArrayRef`] entry point. +//! +//! # Cost model +//! +//! The composition cost is one virtual call per chunk per encoding level (amortized over ~1024 +//! elements, i.e. fractions of a nanosecond per element) plus, for each wrapper level, one +//! in-place pass over an L1-resident chunk. There is no per-element dynamic dispatch and no +//! heap-allocated state introduced on the way down: the sink chain lives in the caller's stack +//! frames, one frame per encoding level. Leaf producers reuse whatever scratch buffer their +//! decompressor already maintains (e.g. the FastLanes 1024-element unpack buffer). +//! +//! Streaming is an explicit capability, not a silent guarantee: encodings advertise it via +//! [`VTable::supports_decompress_chunks`](crate::vtable::VTable::supports_decompress_chunks) +//! (wrappers propagate the check through their children), and +//! [`ArrayRef::decompress_chunks`] errors on an unsupported tree instead of quietly +//! materializing it. Callers that want a works-everywhere path opt into the two-pass fallback +//! by name with [`ArrayRef::decompress_chunks_or_materialize`]. +//! +//! # Contract +//! +//! - Chunks arrive in order, are contiguous, and cover `0..array.len()` exactly. +//! - Chunks are producer-owned scratch: the sink may mutate them freely (wrappers rely on this to +//! transform values in place), and their contents are invalid once `accept` returns. +//! - Validity is *not* streamed: positions that are logically null contain unspecified (but +//! initialized) values, matching `execute`'s decompression behavior. Callers that need null +//! information should fetch `array.validity()` separately. +//! - Currently only primitive-typed arrays are supported. + +use std::marker::PhantomData; +use std::ops::Range; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::PrimitiveArray; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; + +/// The target number of elements per streamed chunk. +/// +/// This matches the FastLanes block size so leaf decompressors can hand their unpack scratch +/// buffer to the sink without copying. Producers may emit shorter chunks (e.g. a sliced first or +/// last block). +pub const DECOMPRESS_CHUNK_LEN: usize = 1024; + +/// A type-erased, mutable view over one chunk of decompressed primitive values. +/// +/// This is a `(PType, *mut, len)` triple rather than a generic slice so it can cross the +/// `dyn ChunkSink` boundary; sinks recover the typed slice with [`Self::as_slice_mut`]. The +/// erasure cost is paid once per chunk, not per element. +pub struct ChunkMut<'a> { + ptype: PType, + data: *mut u8, + len: usize, + _marker: PhantomData<&'a mut u8>, +} + +impl<'a> ChunkMut<'a> { + /// Wrap a typed slice of decompressed values. + pub fn new(values: &'a mut [T]) -> Self { + Self { + ptype: T::PTYPE, + data: values.as_mut_ptr().cast(), + len: values.len(), + _marker: PhantomData, + } + } + + /// The primitive type of the values in this chunk. + #[inline] + pub fn ptype(&self) -> PType { + self.ptype + } + + /// The number of values in this chunk. + #[inline] + #[allow(clippy::len_without_is_empty)] + pub fn len(&self) -> usize { + self.len + } + + /// View the chunk as a typed slice. + /// + /// # Panics + /// Panics if `T::PTYPE` does not match the chunk's ptype. + #[inline] + pub fn as_slice(&self) -> &[T] { + assert_eq!(T::PTYPE, self.ptype, "ChunkMut ptype mismatch"); + // SAFETY: constructed from a valid `&mut [T]` with matching ptype; the lifetime is tied + // to the original borrow via `_marker`. + unsafe { std::slice::from_raw_parts(self.data.cast(), self.len) } + } + + /// View the chunk as a mutable typed slice. + /// + /// # Panics + /// Panics if `T::PTYPE` does not match the chunk's ptype. + #[inline] + pub fn as_slice_mut(&mut self) -> &mut [T] { + assert_eq!(T::PTYPE, self.ptype, "ChunkMut ptype mismatch"); + // SAFETY: constructed from a valid, exclusively borrowed `&mut [T]` with matching ptype. + unsafe { std::slice::from_raw_parts_mut(self.data.cast(), self.len) } + } + + /// Reborrow this chunk with a shorter lifetime, e.g. to forward it to a downstream sink. + #[inline] + pub fn reborrow(&mut self) -> ChunkMut<'_> { + ChunkMut { + ptype: self.ptype, + data: self.data, + len: self.len, + _marker: PhantomData, + } + } +} + +/// Consumer side of [`ArrayRef::decompress_chunks`]. +/// +/// Implementations receive each decompressed chunk exactly once, in array order. `row_range` is +/// the range of logical rows (relative to the array being iterated) that `chunk` covers; its +/// length always equals `chunk.len()`. +pub trait ChunkSink { + /// Accept the next chunk of decompressed values. + fn accept(&mut self, chunk: ChunkMut<'_>, row_range: Range) -> VortexResult<()>; +} + +impl ChunkSink for F +where + F: FnMut(ChunkMut<'_>, Range) -> VortexResult<()>, +{ + #[inline] + fn accept(&mut self, chunk: ChunkMut<'_>, row_range: Range) -> VortexResult<()> { + self(chunk, row_range) + } +} + +impl ArrayRef { + /// Returns whether this array (recursively, through wrapper encodings) can stream its + /// decompressed values via [`Self::decompress_chunks`] without materializing anything. + pub fn supports_decompress_chunks(&self) -> bool { + self.dtype().is_primitive() && self.dyn_array().supports_decompress_chunks(self) + } + + /// Stream the array's decompressed values through `sink` in cache-resident chunks. + /// + /// See the [module docs](self) for the contract and cost model. This errors — without + /// emitting any chunks — if the encoding tree does not support streaming (check with + /// [`Self::supports_decompress_chunks`]); it never silently falls back to full + /// materialization. Use [`Self::decompress_chunks_or_materialize`] to opt into that + /// fallback explicitly. + pub fn decompress_chunks( + &self, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, + ) -> VortexResult<()> { + vortex_ensure!( + self.supports_decompress_chunks(), + "decompress_chunks is not supported by this array tree (root encoding {}, dtype {}); \ + use decompress_chunks_or_materialize for an explicit two-pass fallback", + self.encoding_id(), + self.dtype() + ); + self.decompress_chunks_unchecked(ctx, sink) + } + + /// Stream chunks if the tree supports it, otherwise fall back to executing the array to + /// canonical and streaming the materialized result — the two-pass behavior, chosen by name. + pub fn decompress_chunks_or_materialize( + &self, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, + ) -> VortexResult<()> { + vortex_ensure!( + self.dtype().is_primitive(), + "decompress_chunks requires a primitive-typed array, got {}", + self.dtype() + ); + if self.supports_decompress_chunks() { + self.decompress_chunks_unchecked(ctx, sink) + } else { + decompress_chunks_via_canonical(self, ctx, sink) + } + } + + fn decompress_chunks_unchecked( + &self, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, + ) -> VortexResult<()> { + #[cfg(debug_assertions)] + { + let mut checked = CoverageCheckSink { + inner: sink, + next_row: 0, + ptype: self.dtype().as_ptype(), + }; + self.dyn_array() + .decompress_chunks(self, ctx, &mut checked)?; + debug_assert_eq!( + checked.next_row, + self.len(), + "decompress_chunks did not cover the full array" + ); + Ok(()) + } + #[cfg(not(debug_assertions))] + self.dyn_array().decompress_chunks(self, ctx, sink) + } +} + +#[cfg(debug_assertions)] +struct CoverageCheckSink<'a> { + inner: &'a mut dyn ChunkSink, + next_row: usize, + ptype: PType, +} + +#[cfg(debug_assertions)] +impl ChunkSink for CoverageCheckSink<'_> { + fn accept(&mut self, chunk: ChunkMut<'_>, row_range: Range) -> VortexResult<()> { + debug_assert_eq!(row_range.start, self.next_row, "non-contiguous chunk"); + debug_assert_eq!( + row_range.len(), + chunk.len(), + "chunk/row_range length mismatch" + ); + debug_assert_eq!(chunk.ptype(), self.ptype, "chunk ptype mismatch"); + self.next_row = row_range.end; + self.inner.accept(chunk, row_range) + } +} + +/// Fallback chunked decompression: execute the array to a canonical [`PrimitiveArray`], then +/// stream copies of its values in [`DECOMPRESS_CHUNK_LEN`]-sized chunks. +/// +/// This is the two-pass baseline. It copies each chunk into a single reusable scratch buffer +/// (allocated once) because sinks receive exclusive, mutable chunks. +pub fn decompress_chunks_via_canonical( + array: &ArrayRef, + ctx: &mut ExecutionCtx, + sink: &mut dyn ChunkSink, +) -> VortexResult<()> { + let primitive = array.clone().execute::(ctx)?; + match_each_native_ptype!(primitive.ptype(), |T| { + stream_slice_chunks::(primitive.as_slice::(), sink) + }) +} + +fn stream_slice_chunks(values: &[T], sink: &mut dyn ChunkSink) -> VortexResult<()> { + let mut scratch = vec![T::default(); values.len().min(DECOMPRESS_CHUNK_LEN)]; + for (chunk_idx, chunk) in values.chunks(DECOMPRESS_CHUNK_LEN).enumerate() { + let start = chunk_idx * DECOMPRESS_CHUNK_LEN; + let scratch = &mut scratch[..chunk.len()]; + scratch.copy_from_slice(chunk); + sink.accept(ChunkMut::new(scratch), start..start + chunk.len())?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + + use super::*; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::ConstantArray; + use crate::arrays::Patched; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::patches::Patches; + use crate::scalar::Scalar; + + fn collect_chunks(array: &ArrayRef) -> VortexResult> { + let mut ctx = array_session().create_execution_ctx(); + let mut out = Vec::with_capacity(array.len()); + array.decompress_chunks(&mut ctx, &mut |chunk: ChunkMut<'_>, + _range: Range| + -> VortexResult<()> { + out.extend_from_slice(chunk.as_slice::()); + Ok(()) + })?; + Ok(out) + } + + #[test] + fn constant_chunks() -> VortexResult<()> { + // Length deliberately not a multiple of the chunk size. + let array = ConstantArray::new(7i32, 2500).into_array(); + assert!(array.supports_decompress_chunks()); + let chunked = collect_chunks::(&array)?; + assert_eq!(chunked, vec![7i32; 2500]); + Ok(()) + } + + #[test] + fn null_constant_chunks_cover_length() -> VortexResult<()> { + let array = ConstantArray::new( + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + 100, + ) + .into_array(); + // Values are unspecified for nulls; only coverage matters (checked in debug builds too). + let chunked = collect_chunks::(&array)?; + assert_eq!(chunked.len(), 100); + Ok(()) + } + + #[test] + fn patched_over_constant_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let inner = ConstantArray::new(0u16, 2500).into_array(); + let patches = Patches::new( + 2500, + 0, + buffer![1u32, 1023, 1024, 2047, 2499].into_array(), + buffer![11u16, 22, 33, 44, 55].into_array(), + None, + )?; + let array = Patched::from_array_and_patches(inner, &patches, &mut ctx)?.into_array(); + + assert!(array.supports_decompress_chunks()); + let chunked = collect_chunks::(&array)?; + let expected = array.execute::(&mut ctx)?; + assert_eq!(chunked.as_slice(), expected.as_slice::()); + Ok(()) + } + + #[test] + fn unsupported_encoding_errors_without_emitting() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // A primitive-typed encoding tree with no streaming support: Dict over primitives. + let array = crate::arrays::DictArray::try_new( + buffer![0u32, 1, 0, 1].into_array(), + buffer![10i32, 20].into_array(), + )? + .into_array(); + assert!(!array.supports_decompress_chunks()); + + let mut emitted = 0usize; + let result = array.decompress_chunks(&mut ctx, &mut |chunk: ChunkMut<'_>, + _range: Range| + -> VortexResult<()> { + emitted += chunk.len(); + Ok(()) + }); + assert!(result.is_err()); + assert_eq!(emitted, 0, "no chunks may be emitted on unsupported trees"); + + // The explicit fallback still works and covers the array. + let mut out = Vec::new(); + array.decompress_chunks_or_materialize(&mut ctx, &mut |chunk: ChunkMut<'_>, + _range: Range| + -> VortexResult<()> { + out.extend_from_slice(chunk.as_slice::()); + Ok(()) + })?; + assert_eq!(out, vec![10i32, 20, 10, 20]); + Ok(()) + } +} diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 9439dc7dd1b..4d085b22e91 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -113,6 +113,7 @@ pub mod buffer; pub mod builders; pub mod builtins; mod canonical; +pub mod chunk_iter; mod columnar; pub mod compute; pub mod display;