From 93b376cde74efd1c0ef20eed224c03bd01f47601 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 25 Aug 2026 09:00:08 +0000 Subject: [PATCH 1/5] Add composable chunked decompression via a vtable method Introduce VTable::decompress_chunks: a push-based streaming decompression API that walks an array's decompressed values in cache-resident ~1024 element chunks without materializing the full array. Chunks flow through a ChunkSink chain: leaf encodings stream their existing unpack scratch buffers, and wrapper encodings compose by interposing a stack-allocated sink adapter and recursing into their child through the erased ArrayRef entry point, so the mechanism composes across arbitrary encoding trees with no per-element dynamic dispatch and no heap state on the way down. The default implementation executes to canonical and streams the result (the two-pass baseline), so the method is always available. BitPacked overrides it to stream FastLanes blocks straight from its unpack scratch (applying patches per block via a cursor), and FoR overrides it by adding the reference value in place per chunk before forwarding. On 4Mi u32 FoR-over-BitPacked, a streaming sum through the vtable path runs ~1.7x faster than fused decompress-then-iterate, within ~22% of a fully monomorphized loop doing identical per-chunk work; the pure dynamic dispatch cost measures ~0.06ns per element. Signed-off-by: "Joe Isaacs" Signed-off-by: Joe Isaacs --- encodings/fastlanes/Cargo.toml | 4 + .../fastlanes/benches/chunked_decompress.rs | 227 +++++++++++++++++ .../bitpacking/array/chunked_decompress.rs | 191 ++++++++++++++ .../fastlanes/src/bitpacking/array/mod.rs | 1 + .../fastlanes/src/bitpacking/vtable/mod.rs | 10 + .../fastlanes/src/for/array/for_decompress.rs | 47 ++++ encodings/fastlanes/src/for/vtable/mod.rs | 10 + vortex-array/src/array/mod.rs | 18 ++ vortex-array/src/array/vtable/mod.rs | 20 ++ vortex-array/src/chunk_iter.rs | 232 ++++++++++++++++++ vortex-array/src/lib.rs | 1 + 11 files changed, 761 insertions(+) create mode 100644 encodings/fastlanes/benches/chunked_decompress.rs create mode 100644 encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs create mode 100644 vortex-array/src/chunk_iter.rs 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..73c28bf1c2d --- /dev/null +++ b/encodings/fastlanes/benches/chunked_decompress.rs @@ -0,0 +1,227 @@ +// 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: BitPacked streams unpacked blocks, FoR shifts each +/// block in place, the sink folds it — one streaming pass, no materialization. +#[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) + }); +} + +/// 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()) + }); +} 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..b9865e97fc2 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs @@ -0,0 +1,191 @@ +// 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::match_each_integer_ptype; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_error::VortexResult; + +use crate::BitPacked; +use crate::BitPackedArrayExt; +use crate::unpack_iter::BitPacked as BitPackedUnpack; + +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(()); + } + + // Patches are sparse: materialize them once as (local row, value) pairs so the per-chunk + // loop only advances a cursor. This is the only heap state the streaming path allocates, + // and it is proportional to the patch count, not the array length. + let patch_list: Vec<(usize, T)> = match array.patches() { + None => Vec::new(), + Some(patches) => { + let indices = patches.indices().clone().execute::(ctx)?; + let values = patches.values().clone().execute::(ctx)?; + let values = values.as_slice::(); + let offset = patches.offset(); + match_each_unsigned_integer_ptype!(indices.ptype(), |P| { + indices + .as_slice::

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

>::as_(idx) - offset, v)) + .collect() + }) + } + }; + + let mut chunks = array.unpacked_chunks::()?; + 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(&mut ctx, &mut |chunk: ChunkMut<'_>, + _range: std::ops::Range| + -> 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)?; + assert_chunks_match_execute::(bp.into_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()) + } + + #[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/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 68fbf1b41d3..4fab39cc63e 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::array::chunked_decompress; use crate::bitpacking::vtable::rules::RULES; mod kernels; mod operations; @@ -297,6 +299,14 @@ impl VTable for BitPacked { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } + + 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..3cdbec2eed8 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; @@ -137,6 +141,49 @@ pub(crate) fn fused_decompress< Ok(builder.finish_into_primitive()) } +/// Streaming chunked decompression for FoR arrays. +/// +/// FoR composes 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<()> { + 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) + } + }) +} + +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..08633749a14 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,14 @@ impl VTable for FoR { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(decompress(&array, ctx)?.into_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..a2c6c5cc0f5 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -220,6 +220,14 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { ctx: &mut ExecutionCtx, ) -> VortexResult; + /// 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 +498,16 @@ impl DynArrayData for ArrayData { V::execute(typed, ctx) } + 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..cd4c89db8fe 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -187,6 +187,26 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { /// incorrectly contains null values. fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult; + /// 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 default + /// implementation executes the array to canonical and streams the result (the two-pass + /// baseline), so this is always possible; encodings should override it to stream chunks + /// directly out of their decompression kernel, and wrapper encodings should compose by + /// interposing a stack-allocated [`ChunkSink`](crate::chunk_iter::ChunkSink) adapter and + /// recursing into their child via + /// [`ArrayRef::decompress_chunks`](crate::ArrayRef::decompress_chunks). + /// + /// 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<()> { + crate::chunk_iter::decompress_chunks_via_canonical(array.array(), ctx, sink) + } + /// 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/chunk_iter.rs b/vortex-array/src/chunk_iter.rs new file mode 100644 index 00000000000..46809eecd94 --- /dev/null +++ b/vortex-array/src/chunk_iter.rs @@ -0,0 +1,232 @@ +// 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). +//! +//! The default implementation executes the array to canonical and then streams the result, which +//! is exactly the two-pass behavior this API exists to avoid — so `decompress_chunks` is always +//! possible on any encoding, and monotonically improves as encodings add overrides. +//! +//! # 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 { + /// Stream the array's decompressed values through `sink` in cache-resident chunks. + /// + /// See the [module docs](self) for the contract and cost model. Encodings without a + /// specialized implementation fall back to full decompression followed by chunked iteration + /// of the result. + pub fn decompress_chunks( + &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() + ); + + #[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(()) +} 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; From 880b31ee809b3a9634670bd967d1bcaddb7022d9 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 25 Aug 2026 10:18:58 +0000 Subject: [PATCH 2/5] Stream FoR-over-BitPacked chunks through the fused unpack kernel When a FoR array's child is BitPacked (and the reference is unsigned, mirroring the dispatch in decompress), FoR::decompress_chunks now streams each FastLanes block through the fused FoRStrategy unchecked_unfor_pack kernel instead of wrapping the child's chunk stream with an in-place add pass. Patch values get the reference applied when the patch cursor list is built. The generic sink-composition path remains for any other child. The block streaming and patch-cursor loop is shared between BitPacked's plain path and the fused FoR path via stream_unpacked_chunks, generic over the UnpackStrategy. On 4Mi u32 FoR-over-BitPacked, a streaming sum through the vtable path now matches the hand-written monomorphized fused loop within noise (1.35ms vs 1.36ms median), i.e. the same speed as the unpack_map-style kernels, while the generic composition (measured via a signed reference) stays available for arbitrary children. Signed-off-by: "Joe Isaacs" Signed-off-by: Joe Isaacs --- .../fastlanes/benches/chunked_decompress.rs | 43 ++++++++++- .../bitpacking/array/chunked_decompress.rs | 75 ++++++++++++++----- encodings/fastlanes/src/bitpacking/mod.rs | 1 + .../fastlanes/src/bitpacking/vtable/mod.rs | 2 +- .../fastlanes/src/for/array/for_decompress.rs | 58 +++++++++++++- 5 files changed, 155 insertions(+), 24 deletions(-) diff --git a/encodings/fastlanes/benches/chunked_decompress.rs b/encodings/fastlanes/benches/chunked_decompress.rs index 73c28bf1c2d..50bc0736391 100644 --- a/encodings/fastlanes/benches/chunked_decompress.rs +++ b/encodings/fastlanes/benches/chunked_decompress.rs @@ -81,8 +81,9 @@ impl ChunkSink for SumSink { } } -/// Sum through the composable vtable path: BitPacked streams unpacked blocks, FoR shifts each -/// block in place, the sink folds it — one streaming pass, no materialization. +/// 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(); @@ -97,6 +98,44 @@ fn chunked_vtable_sum(bencher: Bencher) { }); } +/// 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] diff --git a/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs b/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs index b9865e97fc2..0f74e5cc972 100644 --- a/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs @@ -14,13 +14,18 @@ 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>, @@ -41,28 +46,45 @@ fn decompress_chunks_typed( return Ok(()); } - // Patches are sparse: materialize them once as (local row, value) pairs so the per-chunk - // loop only advances a cursor. This is the only heap state the streaming path allocates, - // and it is proportional to the patch count, not the array length. - let patch_list: Vec<(usize, T)> = match array.patches() { + let patch_list = match array.patches() { None => Vec::new(), - Some(patches) => { - let indices = patches.indices().clone().execute::(ctx)?; - let values = patches.values().clone().execute::(ctx)?; - let values = values.as_slice::(); - let offset = patches.offset(); - match_each_unsigned_integer_ptype!(indices.ptype(), |P| { - indices - .as_slice::

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

>::as_(idx) - offset, v)) - .collect() - }) - } + 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| { @@ -173,6 +195,23 @@ mod tests { 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(); 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 4fab39cc63e..8b0f4d9dd24 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -50,7 +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::array::chunked_decompress; +use crate::bitpacking::chunked_decompress; use crate::bitpacking::vtable::rules::RULES; mod kernels; mod operations; diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index 3cdbec2eed8..36e4fa7dc8b 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -25,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; @@ -143,14 +144,28 @@ pub(crate) fn fused_decompress< /// Streaming chunked decompression for FoR arrays. /// -/// FoR composes 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. +/// 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() @@ -169,6 +184,43 @@ pub(crate) fn decompress_chunks( }) } +/// 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, From 4cfeef185a948ffbb04e5d001c2becad0b8f6cee Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 25 Aug 2026 11:36:41 +0000 Subject: [PATCH 3/5] Add decompress_chunks for Constant and Patched arrays Constant streams by filling a single stack-resident 1024-element scratch chunk with the constant and re-emitting it across the array (refilled per emission, since sinks may mutate chunks in place), so the array is never materialized. Patched streams its inner child through the erased decompress_chunks entry point and interposes a sink adapter that overwrites patched rows in each chunk before forwarding. The lane-transposed patch layout is flattened once up front into row-sorted (row, value) pairs so the per-chunk work is a single cursor advance. This composes with any inner encoding: Patched(Constant) patches each re-emitted constant chunk, Patched(BitPacked) patches each unpacked FastLanes block. On 4Mi u32 Patched-over-Constant with sparse patches, a streaming sum runs 3.6x faster than execute-then-iterate (0.59ms vs 2.12ms median) since the 16MB base buffer is never written. Signed-off-by: "Joe Isaacs" Signed-off-by: Joe Isaacs --- .../fastlanes/benches/chunked_decompress.rs | 54 +++++++++ .../src/arrays/constant/vtable/mod.rs | 31 ++++++ vortex-array/src/arrays/patched/vtable/mod.rs | 103 ++++++++++++++++++ vortex-array/src/chunk_iter.rs | 69 ++++++++++++ 4 files changed, 257 insertions(+) diff --git a/encodings/fastlanes/benches/chunked_decompress.rs b/encodings/fastlanes/benches/chunked_decompress.rs index 50bc0736391..e97ea282eba 100644 --- a/encodings/fastlanes/benches/chunked_decompress.rs +++ b/encodings/fastlanes/benches/chunked_decompress.rs @@ -264,3 +264,57 @@ fn fused_decompress(bencher: Bencher) { 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) + }); +} diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 6917d979239..4ad5f919f21 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,34 @@ impl VTable for Constant { )?)) } + 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..cc442f8fe31 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,107 @@ impl VTable for Patched { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } + + 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/chunk_iter.rs b/vortex-array/src/chunk_iter.rs index 46809eecd94..7081ba954c1 100644 --- a/vortex-array/src/chunk_iter.rs +++ b/vortex-array/src/chunk_iter.rs @@ -230,3 +230,72 @@ fn stream_slice_chunks(values: &[T], sink: &mut dyn ChunkSink) - } 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(); + 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(); + + let chunked = collect_chunks::(&array)?; + let expected = array.execute::(&mut ctx)?; + assert_eq!(chunked.as_slice(), expected.as_slice::()); + Ok(()) + } +} From eb92b00ab5e79c6c75a0a5c2f713c0709db8b42d Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 25 Aug 2026 12:27:02 +0000 Subject: [PATCH 4/5] Benchmark chunked streaming against sparse decompression Add benches comparing the chunked decompress_chunks path on Patched(Constant) against the sparse decompression baseline (execute: canonicalize the constant into a full buffer, then scatter patches), for both pure materialization and materialize-then-sum. Signed-off-by: "Joe Isaacs" Signed-off-by: Joe Isaacs --- .../fastlanes/benches/chunked_decompress.rs | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/encodings/fastlanes/benches/chunked_decompress.rs b/encodings/fastlanes/benches/chunked_decompress.rs index e97ea282eba..2a6867b4c4f 100644 --- a/encodings/fastlanes/benches/chunked_decompress.rs +++ b/encodings/fastlanes/benches/chunked_decompress.rs @@ -318,3 +318,41 @@ fn two_pass_sum_patched_constant(bencher: Bencher) { 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()) + }); +} From f62481a88e1aaba97723a8eb95bea6d8adf3cf74 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 25 Aug 2026 14:24:47 +0000 Subject: [PATCH 5/5] Make chunked decompression an explicit capability, not a silent fallback The default VTable::decompress_chunks previously fell back to executing the array to canonical and streaming the materialized result, which silently defeated the purpose of the API: decompressing and transforming blocks while they are L1-resident. Streaming is now advertised via VTable::supports_decompress_chunks (default false); wrapper encodings propagate the check through the children they stream from, so support of the whole tree is decided before any chunk is emitted. ArrayRef::decompress_chunks errors on unsupported trees without emitting anything, and the materializing two-pass fallback moves to the explicitly named ArrayRef::decompress_chunks_or_materialize. Primitive gains a real streaming implementation (it is already decompressed; chunks are copied through one reusable L1 scratch buffer), Constant and BitPacked advertise support directly, FoR advertises support when its fused BitPacked path applies or its child supports streaming, and Patched advertises support when its inner array does. Signed-off-by: "Joe Isaacs" Signed-off-by: Joe Isaacs --- .../bitpacking/array/chunked_decompress.rs | 10 ++- .../fastlanes/src/bitpacking/vtable/mod.rs | 4 + .../fastlanes/src/for/array/for_decompress.rs | 7 ++ encodings/fastlanes/src/for/vtable/mod.rs | 4 + vortex-array/src/array/mod.rs | 8 ++ vortex-array/src/array/vtable/mod.rs | 40 ++++++--- .../src/arrays/constant/vtable/mod.rs | 4 + vortex-array/src/arrays/patched/vtable/mod.rs | 4 + .../src/arrays/primitive/vtable/mod.rs | 27 ++++++ vortex-array/src/chunk_iter.rs | 86 +++++++++++++++++-- 10 files changed, 175 insertions(+), 19 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs b/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs index 0f74e5cc972..a5ae674ebcc 100644 --- a/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/chunked_decompress.rs @@ -128,8 +128,10 @@ mod tests { 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(&mut ctx, &mut |chunk: ChunkMut<'_>, - _range: std::ops::Range| + array.decompress_chunks_or_materialize(&mut ctx, &mut |chunk: ChunkMut<'_>, + _range: std::ops::Range< + usize, + >| -> VortexResult<()> { out.extend_from_slice(chunk.as_slice::()); Ok(()) @@ -152,7 +154,9 @@ mod tests { 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)?; - assert_chunks_match_execute::(bp.into_array()) + let array = bp.into_array(); + assert!(array.supports_decompress_chunks()); + assert_chunks_match_execute::(array) } #[test] diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 8b0f4d9dd24..995d68b7e0d 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -300,6 +300,10 @@ impl VTable for BitPacked { RULES.evaluate(array, parent, child_idx) } + fn supports_decompress_chunks(_array: ArrayView<'_, Self>) -> bool { + true + } + fn decompress_chunks( array: ArrayView<'_, Self>, ctx: &mut ExecutionCtx, diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index 36e4fa7dc8b..1f68b11bc6b 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -142,6 +142,13 @@ 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 diff --git a/encodings/fastlanes/src/for/vtable/mod.rs b/encodings/fastlanes/src/for/vtable/mod.rs index 08633749a14..0d8187dce5b 100644 --- a/encodings/fastlanes/src/for/vtable/mod.rs +++ b/encodings/fastlanes/src/for/vtable/mod.rs @@ -164,6 +164,10 @@ impl VTable for FoR { 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, diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index a2c6c5cc0f5..9a7d021ead7 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -220,6 +220,9 @@ 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, @@ -498,6 +501,11 @@ 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, diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index cd4c89db8fe..8252f1b1aa7 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -187,24 +187,44 @@ 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 default - /// implementation executes the array to canonical and streams the result (the two-pass - /// baseline), so this is always possible; encodings should override it to stream chunks - /// directly out of their decompression kernel, and wrapper encodings should compose by - /// interposing a stack-allocated [`ChunkSink`](crate::chunk_iter::ChunkSink) adapter and - /// recursing into their child via - /// [`ArrayRef::decompress_chunks`](crate::ArrayRef::decompress_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). /// - /// The caller guarantees the array is primitive-typed. Implementations must emit contiguous, - /// in-order chunks covering exactly `0..len` (checked in debug builds). + /// 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<()> { - crate::chunk_iter::decompress_chunks_via_canonical(array.array(), ctx, sink) + _ = (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. diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 4ad5f919f21..7b97e3a9478 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -189,6 +189,10 @@ impl VTable for Constant { )?)) } + fn supports_decompress_chunks(_array: ArrayView<'_, Self>) -> bool { + true + } + fn decompress_chunks( array: ArrayView<'_, Self>, _ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index cc442f8fe31..a7a45e4424f 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -312,6 +312,10 @@ impl VTable for Patched { 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, 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 index 7081ba954c1..735115b5fa0 100644 --- a/vortex-array/src/chunk_iter.rs +++ b/vortex-array/src/chunk_iter.rs @@ -19,9 +19,12 @@ //! frames, one frame per encoding level. Leaf producers reuse whatever scratch buffer their //! decompressor already maintains (e.g. the FastLanes 1024-element unpack buffer). //! -//! The default implementation executes the array to canonical and then streams the result, which -//! is exactly the two-pass behavior this API exists to avoid — so `decompress_chunks` is always -//! possible on any encoding, and monotonically improves as encodings add overrides. +//! 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 //! @@ -145,22 +148,58 @@ where } 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. Encodings without a - /// specialized implementation fall back to full decompression followed by chunked iteration - /// of the result. + /// 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 { @@ -262,6 +301,7 @@ mod tests { 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(()) @@ -293,9 +333,43 @@ mod tests { )?; 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(()) + } }