From 9fdbe75a3dd0f320873c41eb9eeebbe84eaa4a67 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 30 Jul 2026 22:17:51 +0100 Subject: [PATCH 1/5] Zstd: append values in bulk and harden the frame metadata reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `append_to_builder` walked the decompressed frames one value at a time, appending each through a single-value call and interleaving null runs by hand. It now derives the value byte range from the slice's own metadata and hands the whole region to `append_value_slices`, which sizes the offsets, byte storage and validity once — one offset store plus one `memcpy` per value. Adds a `VarBinViewBuilder` path as well: the frames already hold the values contiguously, so the views can reference them in place instead of going through the canonical array and rewriting every view a second time to rebase its buffer index. `reconstruct_views` therefore takes the buffer index the views should start at. Frame metadata comes straight off disk, so the arithmetic it drives is now checked and surfaced as errors rather than panics or wrapped lengths: value counts, frame sizes, length prefixes and the offsets they walk to. The missing `n_values` fallback no longer reads a byte count as a value count for variable-width values, where that mis-attributes values to frames — it is accepted only for the single-frame case that is still recoverable. Decompression also writes through `WriteBuf` into uninitialized spare capacity rather than a `&mut [u8]` over memory nothing has written yet. Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + encodings/zstd/Cargo.toml | 1 + encodings/zstd/src/array.rs | 622 +++++++++++++++++++---- encodings/zstd/src/test.rs | 142 ++++++ vortex-cuda/src/kernel/encodings/zstd.rs | 2 +- 5 files changed, 669 insertions(+), 99 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6b69e6f253a..f17b7234f4a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10605,6 +10605,7 @@ name = "vortex-zstd" version = "0.1.0" dependencies = [ "itertools 0.14.0", + "num-traits", "prost 0.14.4", "rstest", "vortex-array", diff --git a/encodings/zstd/Cargo.toml b/encodings/zstd/Cargo.toml index 1b2cb4f2dd0..49f15619807 100644 --- a/encodings/zstd/Cargo.toml +++ b/encodings/zstd/Cargo.toml @@ -25,6 +25,7 @@ unstable_encodings = [] [dependencies] itertools = { workspace = true } +num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 41e436a1bea..82b7b3d8a92 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -6,9 +6,12 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; +use std::mem::MaybeUninit; +use std::ops::Range; use std::sync::Arc; use itertools::Itertools as _; +use num_traits::AsPrimitive; use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; @@ -31,6 +34,7 @@ use vortex_array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinBuilder; +use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::OffsetBuilderPType; use vortex_array::match_each_varbin_builder; @@ -54,10 +58,10 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use zstd::zstd_safe::WriteBuf; use crate::ZstdFrameMetadata; use crate::ZstdMetadata; @@ -284,12 +288,14 @@ impl VTable for Zstd { { return result; } - array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx) + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. There is deliberately no + // canonicalize-then-append fallback — it would decompress to a `VarBinView` only for + // `VarBinView::append_to_builder` to reject the same remainder. + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for Zstd requires a variable-binary builder") + }; + append_to_varbinview(array, builder, ctx) } fn reduce_parent( @@ -301,47 +307,97 @@ impl VTable for Zstd { } } -/// Copies the decompressed values into `builder`. +fn unsliced_validity(array: ArrayView<'_, Zstd>) -> Validity { + child_to_validity( + array.slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ) +} + +/// Copies the decompressed values straight into `builder`'s byte storage. +/// +/// The decompressed frames interleave a length prefix with each value, so the bytes have to be +/// compacted; sizing the offsets, byte storage and validity from the slice's own metadata keeps +/// that down to one offset store plus one `memcpy` per value. fn append_to_varbin( array: ArrayView<'_, Zstd>, builder: &mut VarBinBuilder, ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let slice = array + .data() + .decompress_slice(array.dtype(), &unsliced_validity(array), ctx)?; + let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; + let (values, num_bytes) = slice.value_bytes()?; + // Each value is length-prefixed, so the frames can only be walked in order — which is the + // order `append_valid_slices` visits the valid rows in. A stream that runs out early yields + // empty slices and so fails the builder's byte-count check. + let mut values = zstd_values(values); + builder.append_valid_slices(num_bytes, &mask, |_| values.next().unwrap_or_default()) +} + +/// Hands the decompressed frames to `builder` as data buffers with views built over them. +/// +/// The frames already hold the values contiguously, so the views can reference them in place and +/// the only per-row work is building one view; going through the canonical array instead would +/// rewrite every view a second time to rebase its buffer index. +fn append_to_varbinview( + array: ArrayView<'_, Zstd>, + builder: &mut VarBinViewBuilder, + ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let unsliced_validity = child_to_validity( - array.slots()[ZstdSlots::VALIDITY].as_ref(), - array.dtype().nullability(), - ); let slice = array .data() - .decompress_slice(array.dtype(), &unsliced_validity, ctx)?; - let value_start = slice.value_idx_start - slice.n_skipped_values; - let value_count = slice.value_idx_stop - slice.value_idx_start; - let mut values = zstd_values(slice.bytes.as_slice()) - .skip(value_start) - .take(value_count); + .decompress_slice(array.dtype(), &unsliced_validity(array), ctx)?; let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; - match mask.indices() { - AllOr::All => { - for value in values { - builder.append_n_values(value, 1)?; - } - } - AllOr::None => builder.push_nulls(slice.n_rows), - AllOr::Some(valid_indices) => { - let mut row = 0; - for &valid_index in valid_indices { - builder.push_nulls(valid_index - row); - builder.append_n_values( - values - .next() - .vortex_expect("Zstd value count must match validity"), - 1, - )?; - row = valid_index + 1; - } - builder.push_nulls(slice.n_rows - row); - } + + // No values were stored, so there is nothing to reference and the frames can be dropped. + if mask.all_false() { + builder.append_nulls(slice.n_rows); + return Ok(()); } + + // Build the views against the index the pushed buffers will land at, so the builder does not + // have to rebase them afterwards. + let next_buffer_index = builder.completed_block_count() + u32::from(builder.in_progress()); + // The decompressed frames cover whole frames and so can extend past the requested values on + // either side. Reconstructing over just the requested region keeps the pushed buffers fully + // utilized, which is what `push_buffer_and_adjusted_views` requires: it hands them to the + // finished array as they are, without compacting them. + let value_bytes = slice.bytes.slice(slice.value_byte_range()?); + let (buffers, valid_views) = + try_reconstruct_views(&value_bytes, next_buffer_index, MAX_BUFFER_LEN)?; + vortex_ensure!( + valid_views.len() == mask.true_count(), + "Corrupt zstd metadata: the decompressed frames hold {} values for the {} valid rows of \ + the slice", + valid_views.len(), + mask.true_count() + ); + + let views = match mask.bit_buffer() { + AllOr::All => valid_views, + AllOr::None => unreachable!("handled above"), + AllOr::Some(bits) => { + // Null rows carry an empty view, so scatter the stored values into their rows. Walking + // the set bits a word at a time avoids materializing the mask's indices, which the + // views are the only consumer of. + let mut views = BufferMut::::zeroed(slice.n_rows); + let mut valid_row = 0; + bits.for_each_set_index(|index| { + // In bounds: `valid_views.len() == mask.true_count()` was checked above, and + // `index < slice.n_rows` because `bits` is the mask over those rows. + views[index] = valid_views[valid_row]; + valid_row += 1; + }); + views.freeze() + } + }; + + builder.push_buffer_and_adjusted_views(&buffers, &views, mask); Ok(()) } @@ -545,22 +601,48 @@ fn collect_valid_vbv( /// tests to exercise the splitting path without allocating >2 GiB. pub fn reconstruct_views( buffer: &ByteBuffer, + start_buf_index: u32, max_buffer_len: usize, ) -> (Vec, Buffer) { + let (buffers, views, _) = walk_views(buffer, start_buf_index, max_buffer_len); + (buffers, views) +} + +/// [`reconstruct_views`], but rejecting a buffer that the length prefixes do not tile exactly. +fn try_reconstruct_views( + buffer: &ByteBuffer, + start_buf_index: u32, + max_buffer_len: usize, +) -> VortexResult<(Vec, Buffer)> { + match walk_views(buffer, start_buf_index, max_buffer_len) { + (buffers, views, None) => Ok((buffers, views)), + (_, _, Some(error)) => Err(error), + } +} + +/// Walks `buffer` until it is exhausted or a length prefix leaves it, returning what was decoded +/// along with the error that stopped the walk. +fn walk_views( + buffer: &ByteBuffer, + start_buf_index: u32, + max_buffer_len: usize, +) -> (Vec, Buffer, Option) { let mut views = BufferMut::::empty(); let mut buffers = Vec::new(); let mut segment_start: usize = 0; let mut offset = 0; + // Only a new segment changes the buffer index, so it is tracked instead of recomputed per view. + let mut buf_index = start_buf_index; + let mut error = None; while offset < buffer.len() { - let str_len = ViewLen::from_le_bytes( - buffer - .get(offset..offset + size_of::()) - .vortex_expect("corrupted zstd length") - .try_into() - .ok() - .vortex_expect("must fit ViewLen size"), - ) as usize; + let str_len = match zstd_value_len(buffer.as_slice(), offset) { + Ok(str_len) => str_len, + Err(err) => { + error = Some(err); + break; + } + }; let value_data_offset = offset + size_of::(); let local_offset = value_data_offset - segment_start; @@ -568,12 +650,28 @@ pub fn reconstruct_views( if local_offset + str_len > max_buffer_len && offset > segment_start { buffers.push(buffer.slice(segment_start..offset)); segment_start = offset; + let Some(next_index) = buf_index.checked_add(1) else { + error = Some(vortex_err!("Zstd values need more than u32::MAX buffers")); + break; + }; + buf_index = next_index; } - let local_offset = u32::try_from(value_data_offset - segment_start) - .vortex_expect("local offset within segment must fit in u32"); - let buf_index = u32::try_from(buffers.len()).vortex_expect("buffer index must fit in u32"); - let value = &buffer[value_data_offset..value_data_offset + str_len]; + let Ok(local_offset) = u32::try_from(value_data_offset - segment_start) else { + error = Some(vortex_err!( + "Zstd value offset {} does not fit in u32; max_buffer_len {max_buffer_len} is too large", + value_data_offset - segment_start + )); + break; + }; + let Some(value) = buffer.get(value_data_offset..value_data_offset + str_len) else { + error = Some(vortex_err!( + "Corrupt zstd value: {str_len} bytes at offset {value_data_offset} run past the \ + end of the {} byte frame buffer", + buffer.len() + )); + break; + }; views.push(BinaryView::make_view(value, buf_index, local_offset)); offset = value_data_offset + str_len; } @@ -582,7 +680,63 @@ pub fn reconstruct_views( buffers.push(buffer.slice(segment_start..buffer.len())); } - (buffers, views.freeze()) + (buffers, views.freeze(), error) +} + +/// Narrows the views decoded from the frames down to the values a slice requests. +fn slice_views( + views: &Buffer, + range: Range, +) -> VortexResult> { + vortex_ensure!( + range.end <= views.len(), + "Corrupt zstd metadata: values {}..{} are out of bounds of the {} values held by the \ + decompressed frames", + range.start, + range.end, + views.len() + ); + Ok(views.slice(range)) +} + +/// A zstd output buffer over uninitialized spare capacity. +/// +/// `decompress_to_buffer` writes through a raw pointer and reports how many bytes it produced, so +/// it never reads its destination — but handing it a `&mut [u8]` covering uninitialized memory +/// would be undefined behaviour regardless of what it does with it. [`WriteBuf`] is the interface +/// zstd provides for exactly this case, and it keeps the alternative (zeroing the whole buffer +/// before every decompression) off the hot path. +struct UninitDestination<'a> { + spare: &'a mut [MaybeUninit], + filled: usize, +} + +impl<'a> UninitDestination<'a> { + fn new(spare: &'a mut [MaybeUninit]) -> Self { + Self { spare, filled: 0 } + } +} + +// SAFETY: `as_mut_ptr` and `capacity` describe the whole spare region, so zstd only ever writes +// within it, and `filled_until` merely records the count it reports. `as_slice` is bounded by that +// count, so it never exposes a byte zstd did not write. +unsafe impl WriteBuf for UninitDestination<'_> { + fn as_slice(&self) -> &[u8] { + // SAFETY: zstd reported writing `filled` bytes from the start of `spare`. + unsafe { std::slice::from_raw_parts(self.spare.as_ptr().cast::(), self.filled) } + } + + fn capacity(&self) -> usize { + self.spare.len() + } + + fn as_mut_ptr(&mut self) -> *mut u8 { + self.spare.as_mut_ptr().cast::() + } + + unsafe fn filled_until(&mut self, n: usize) { + self.filled = n; + } } struct DecompressedSlice { @@ -593,23 +747,145 @@ struct DecompressedSlice { value_idx_start: usize, value_idx_stop: usize, n_skipped_values: usize, + /// Number of stored values held by `bytes`, which covers whole frames and so may extend past + /// the requested slice on either side. + n_buffered_values: usize, +} + +impl DecompressedSlice { + /// The range of stored values this slice requests, as an index into `bytes`. + /// + /// The frame metadata that drives `n_skipped_values` is untrusted, so a frame claiming to hold + /// values that precede the ones we decompressed is rejected instead of wrapping. + fn value_range(&self) -> VortexResult> { + let start = self + .value_idx_start + .checked_sub(self.n_skipped_values) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: skipped frames hold {} values, past the first \ + requested value {}", + self.n_skipped_values, + self.value_idx_start + ) + })?; + let end = self + .value_idx_stop + .checked_sub(self.n_skipped_values) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: skipped frames hold {} values, past the last \ + requested value {}", + self.n_skipped_values, + self.value_idx_stop + ) + })?; + vortex_ensure!( + start <= end, + "Corrupt zstd metadata: value range {start}..{end} is not ascending" + ); + Ok(start..end) + } + + /// The bounds within `bytes` of the length-prefixed region holding exactly this slice's values. + /// + /// Walking the length prefixes is a dependent load chain, so both ends are derived from the + /// slice metadata where possible: an unsliced array skips both walks. + fn value_byte_range(&self) -> VortexResult> { + let Range { start, end } = self.value_range()?; + let buffer = self.bytes.as_slice(); + let from = zstd_value_offset(buffer, 0, start)?; + let to = if end == self.n_buffered_values { + buffer.len() + } else { + zstd_value_offset(buffer, from, end - start)? + }; + vortex_ensure!( + from <= to && to <= buffer.len(), + "Corrupt zstd metadata: values {from}..{to} are out of bounds of the {} byte frame \ + buffer", + buffer.len() + ); + Ok(from..to) + } + + /// The length-prefixed region of `bytes` holding exactly this slice's values, along with the + /// total size of those values once the length prefixes are dropped. + fn value_bytes(&self) -> VortexResult<(&[u8], usize)> { + let range = self.value_byte_range()?; + let n_values = self.value_range()?.len(); + let buffer = self.bytes.as_slice(); + let bytes = buffer.get(range.clone()).ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: values {}..{} are out of bounds of the {} byte frame \ + buffer", + range.start, + range.end, + buffer.len() + ) + })?; + // Every value carries a length prefix, so the region must be at least that large. + let prefix_bytes = n_values.checked_mul(size_of::()).ok_or_else(|| { + vortex_err!("Corrupt zstd metadata: value count {n_values} overflows a byte count") + })?; + let num_bytes = bytes.len().checked_sub(prefix_bytes).ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: {n_values} values do not fit in the {} bytes holding them", + bytes.len() + ) + })?; + Ok((bytes, num_bytes)) + } } +/// Returns the byte offset `count` length-prefixed values past `offset`. +/// +/// Each step is bounds-checked by the next prefix read, so only the offset the walk lands on needs +/// a check of its own. +fn zstd_value_offset(buffer: &[u8], mut offset: usize, count: usize) -> VortexResult { + for _ in 0..count { + offset += size_of::() + zstd_value_len(buffer, offset)?; + } + vortex_ensure!( + offset <= buffer.len(), + "Corrupt zstd values: walking {count} values ended at offset {offset}, past the end of \ + the {} byte frame buffer", + buffer.len() + ); + Ok(offset) +} + +/// Reads the length prefix of the value starting at `offset`. +#[inline] +fn zstd_value_len(buffer: &[u8], offset: usize) -> VortexResult { + let prefix = buffer + .get(offset..) + .and_then(|rest| rest.first_chunk::<{ size_of::() }>()) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd values: length prefix at offset {offset} runs past the end of the \ + {} byte frame buffer", + buffer.len() + ) + })?; + Ok(ViewLen::from_le_bytes(*prefix) as usize) +} + +/// Iterates the values of a length-prefixed region, stopping at the first prefix that leaves it. +/// +/// Stopping short leaves the value count and byte total below what the caller declared, which the +/// consuming builder rejects, so the walk does not need to report the error itself. fn zstd_values(buffer: &[u8]) -> impl Iterator { let mut offset = 0; std::iter::from_fn(move || { - if offset == buffer.len() { + if offset >= buffer.len() { return None; } - let len = ViewLen::from_le_bytes( - buffer[offset..offset + size_of::()] - .try_into() - .ok() - .vortex_expect("must fit ViewLen size"), - ) as usize; + let len = zstd_value_len(buffer, offset).ok()?; let value_start = offset + size_of::(); + let value = buffer.get(value_start..value_start + len)?; offset = value_start + len; - Some(&buffer[value_start..offset]) + Some(value) }) } @@ -1011,36 +1287,66 @@ impl ZstdData { // what row offset into the first such frame. let byte_width = Self::byte_width(dtype); let slice_n_rows = self.slice_stop - self.slice_start; - let slice_value_indices = unsliced_validity - .execute_mask(self.unsliced_n_rows, ctx)? - .valid_counts_for_indices(&[self.slice_start, self.slice_stop]); + let unsliced_mask = unsliced_validity.execute_mask(self.unsliced_n_rows, ctx)?; + let slice_value_indices = + unsliced_mask.valid_counts_for_indices(&[self.slice_start, self.slice_stop]); let slice_value_idx_start = slice_value_indices[0]; let slice_value_idx_stop = slice_value_indices[1]; let mut frames_to_decompress = vec![]; let mut value_idx_start = 0; - let mut uncompressed_size_to_decompress = 0; + let mut uncompressed_size_to_decompress = 0usize; let mut n_skipped_values = 0; + let mut n_buffered_values = 0; for (frame, frame_meta) in self.frames.iter().zip(&self.metadata.frames) { if value_idx_start >= slice_value_idx_stop { break; } - let frame_uncompressed_size = usize::try_from(frame_meta.uncompressed_size) - .vortex_expect("Uncompressed size must fit in usize"); - let frame_n_values = if frame_meta.n_values == 0 { - // possibly older primitive-only metadata that just didn't store this + let frame_uncompressed_size = + usize::try_from(frame_meta.uncompressed_size).map_err(|_| { + vortex_err!( + "Zstd frame uncompressed size {} does not fit in a usize", + frame_meta.uncompressed_size + ) + })?; + let frame_n_values = if frame_meta.n_values != 0 { + usize::try_from(frame_meta.n_values).map_err(|_| { + vortex_err!( + "Zstd frame value count {} does not fit in a usize", + frame_meta.n_values + ) + })? + } else if dtype.is_primitive() { + // Possibly older primitive-only metadata that just didn't store this. Fixed-width + // values make the byte count an exact value count. frame_uncompressed_size / byte_width } else { - usize::try_from(frame_meta.n_values).vortex_expect("frame size must fit usize") + // The same fallback would read a byte count as a value count for variable-width + // values, which mis-attributes values to frames. A single frame holds every stored + // value, so that case is still recoverable; anything else is not. + vortex_ensure!( + self.frames.len() == 1, + "Zstd frame metadata for a variable-width array is missing its value count" + ); + unsliced_mask.true_count() }; - let value_idx_stop = value_idx_start + frame_n_values; + // Bounding the running total also bounds the two accumulators below, which partition + // it between the frames we keep and the ones we skip. + let value_idx_stop = value_idx_start.checked_add(frame_n_values).ok_or_else(|| { + vortex_err!("Corrupt zstd metadata: frame value counts overflow a usize") + })?; if value_idx_stop > slice_value_idx_start { // we need this frame frames_to_decompress.push(frame); - uncompressed_size_to_decompress += frame_uncompressed_size; + uncompressed_size_to_decompress = uncompressed_size_to_decompress + .checked_add(frame_uncompressed_size) + .ok_or_else(|| { + vortex_err!("Corrupt zstd metadata: frame sizes overflow a usize") + })?; + n_buffered_values += frame_n_values; } else { n_skipped_values += frame_n_values; } @@ -1057,24 +1363,28 @@ impl ZstdData { uncompressed_size_to_decompress, Alignment::new(byte_width), ); - unsafe { - // safety: we immediately fill all bytes in the following loop, - // assuming our metadata's uncompressed size is correct - decompressed.set_len(uncompressed_size_to_decompress); - } let mut uncompressed_start = 0; for frame in frames_to_decompress { - let uncompressed_written = decompressor - .decompress_to_buffer(frame.as_slice(), &mut decompressed[uncompressed_start..])?; - uncompressed_start += uncompressed_written; + // Decompress straight into the spare capacity. Each frame gets only the region after + // the ones before it, bounded by the size the metadata declared, so a frame that + // expands further than advertised is refused by zstd rather than overrunning. + let mut destination = UninitDestination::new( + &mut decompressed.spare_capacity_mut() + [uncompressed_start..uncompressed_size_to_decompress], + ); + uncompressed_start += + decompressor.decompress_to_buffer(frame.as_slice(), &mut destination)?; } if uncompressed_start != uncompressed_size_to_decompress { - vortex_panic!( + vortex_bail!( "Zstd metadata or frames were corrupt; expected {} bytes but decompressed {}", uncompressed_size_to_decompress, uncompressed_start ); } + // SAFETY: the loop above decompressed exactly `uncompressed_start` bytes into the front of + // the spare capacity, and the check above pins that to the requested length. + unsafe { decompressed.set_len(uncompressed_start) }; let decompressed = decompressed.freeze(); // Last, we slice the exact values requested out of the decompressed data. @@ -1089,7 +1399,7 @@ impl ZstdData { // We ensure that the validity of the decompressed array ALWAYS matches the validity // implied by the DType. if !dtype.is_nullable() && !matches!(slice_validity, Validity::NonNullable) { - assert!( + vortex_ensure!( matches!(slice_validity, Validity::AllValid), "ZSTD array expects to be non-nullable but there are nulls after decompression" ); @@ -1109,6 +1419,7 @@ impl ZstdData { value_idx_start: slice_value_idx_start, value_idx_stop: slice_value_idx_stop, n_skipped_values, + n_buffered_values, }) } @@ -1121,10 +1432,21 @@ impl ZstdData { let slice = self.decompress_slice(dtype, unsliced_validity, ctx)?; match dtype { DType::Primitive(..) => { - let slice_values_buffer = slice.bytes.slice( - (slice.value_idx_start - slice.n_skipped_values) * slice.byte_width - ..(slice.value_idx_stop - slice.n_skipped_values) * slice.byte_width, - ); + let Range { start, end } = slice.value_range()?; + let byte_range = start + .checked_mul(slice.byte_width) + .zip(end.checked_mul(slice.byte_width)) + .filter(|(_, byte_stop)| *byte_stop <= slice.bytes.len()) + .map(|(byte_start, byte_stop)| byte_start..byte_stop) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: values {start}..{end} of {} bytes each are \ + out of bounds of the {} byte frame buffer", + slice.byte_width, + slice.bytes.len() + ) + })?; + let slice_values_buffer = slice.bytes.slice(byte_range); let primitive = PrimitiveArray::from_values_byte_buffer( slice_values_buffer, dtype.as_ptype(), @@ -1138,11 +1460,9 @@ impl ZstdData { DType::Binary(_) | DType::Utf8(_) => { match slice.validity.execute_mask(slice.n_rows, ctx)?.indices() { AllOr::All => { - let (buffers, all_views) = reconstruct_views(&slice.bytes, MAX_BUFFER_LEN); - let valid_views = all_views.slice( - slice.value_idx_start - slice.n_skipped_values - ..slice.value_idx_stop - slice.n_skipped_values, - ); + let (buffers, all_views) = + try_reconstruct_views(&slice.bytes, 0, MAX_BUFFER_LEN)?; + let valid_views = slice_views(&all_views, slice.value_range()?)?; // SAFETY: we properly construct the views inside `reconstruct_views` Ok(unsafe { @@ -1161,11 +1481,9 @@ impl ZstdData { ) .into_array()), AllOr::Some(valid_indices) => { - let (buffers, all_views) = reconstruct_views(&slice.bytes, MAX_BUFFER_LEN); - let valid_views = all_views.slice( - slice.value_idx_start - slice.n_skipped_values - ..slice.value_idx_stop - slice.n_skipped_values, - ); + let (buffers, all_views) = + try_reconstruct_views(&slice.bytes, 0, MAX_BUFFER_LEN)?; + let valid_views = slice_views(&all_views, slice.value_range()?)?; let mut views = BufferMut::::zeroed(slice.n_rows); for (view, index) in valid_views.into_iter().zip_eq(valid_indices) { @@ -1185,7 +1503,7 @@ impl ZstdData { } } } - _ => vortex_panic!("Unsupported dtype for Zstd array: {}", dtype), + _ => vortex_bail!("Unsupported dtype for Zstd array: {}", dtype), } } @@ -1257,9 +1575,17 @@ impl OperationsVTable for Zstd { #[cfg(test)] #[expect(clippy::cast_possible_truncation)] mod tests { + use rstest::rstest; + use vortex_array::validity::Validity; use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + use super::DecompressedSlice; + use super::ViewLen; use super::reconstruct_views; + use super::try_reconstruct_views; + use super::zstd_value_len; + use super::zstd_value_offset; use crate::array::BinaryView; /// Build a Zstd-style interleaved buffer: [u32-LE length][string bytes] repeated. @@ -1273,11 +1599,31 @@ mod tests { ByteBuffer::copy_from(buf.as_slice()) } + /// A slice over `bytes` that requests `value_idx_start..value_idx_stop`. + fn decompressed_slice( + bytes: ByteBuffer, + value_idx_start: usize, + value_idx_stop: usize, + n_skipped_values: usize, + n_buffered_values: usize, + ) -> DecompressedSlice { + DecompressedSlice { + bytes, + validity: Validity::NonNullable, + byte_width: 1, + n_rows: value_idx_stop - value_idx_start, + value_idx_start, + value_idx_stop, + n_skipped_values, + n_buffered_values, + } + } + #[test] fn test_reconstruct_views_no_split() { let strings: &[&[u8]] = &[b"hello", b"world"]; let buf = make_interleaved(strings); - let (buffers, views) = reconstruct_views(&buf, 1024); + let (buffers, views) = reconstruct_views(&buf, 0, 1024); assert_eq!(buffers.len(), 1); assert_eq!(views.len(), 2); @@ -1294,7 +1640,7 @@ mod tests { // so it rolls into a second segment. let strings: &[&[u8]] = &[b"aaaaaaaaaaaaa", b"bbbbbbbbbbbbb"]; let buf = make_interleaved(strings); - let (buffers, views) = reconstruct_views(&buf, 20); + let (buffers, views) = reconstruct_views(&buf, 0, 20); assert_eq!(buffers.len(), 2); assert_eq!(views.len(), 2); @@ -1302,4 +1648,84 @@ mod tests { // Second entry starts a new segment at byte 17 (the length prefix), so local offset = 4. assert_eq!(views[1], BinaryView::make_view(b"bbbbbbbbbbbbb", 1, 4)); } + + /// A buffer whose last entry claims more bytes than remain, as corrupt frame data would. + fn make_overrunning() -> ByteBuffer { + let mut buf = Vec::new(); + buf.extend_from_slice(&5u32.to_le_bytes()); + buf.extend_from_slice(b"hello"); + buf.extend_from_slice(&9u32.to_le_bytes()); + buf.extend_from_slice(b"ab"); + ByteBuffer::copy_from(buf.as_slice()) + } + + #[test] + fn test_reconstruct_views_rejects_overrunning_value() { + let buf = make_overrunning(); + assert!(try_reconstruct_views(&buf, 0, 1024).is_err()); + + // The lenient walk keeps the decodable prefix instead of panicking. + let (buffers, views) = reconstruct_views(&buf, 0, 1024); + assert_eq!(buffers.len(), 1); + assert_eq!(views.len(), 1); + assert_eq!(views[0], BinaryView::make_view(b"hello", 0, 4)); + } + + #[test] + fn test_reconstruct_views_rejects_truncated_length_prefix() { + // A trailing partial length prefix cannot start a value. + let buf = + ByteBuffer::copy_from([5u8, 0, 0, 0, b'h', b'e', b'l', b'l', b'o', 1, 0].as_ref()); + assert!(try_reconstruct_views(&buf, 0, 1024).is_err()); + assert_eq!(reconstruct_views(&buf, 0, 1024).1.len(), 1); + } + + #[rstest] + #[case::truncated_buffer(&[0u8, 0, 0], 0)] + #[case::truncated_tail(&[4u8, 0, 0, 0], 2)] + #[case::offset_at_end(&[4u8, 0, 0, 0], 4)] + #[case::offset_past_end(&[4u8, 0, 0, 0], 64)] + fn test_zstd_value_len_rejects_out_of_bounds(#[case] buffer: &[u8], #[case] offset: usize) { + assert!(zstd_value_len(buffer, offset).is_err()); + } + + #[test] + fn test_zstd_value_offset_rejects_walking_past_the_end() -> VortexResult<()> { + let buf = make_interleaved(&[b"hello", b"world"]); + assert_eq!(zstd_value_offset(buf.as_slice(), 0, 2)?, buf.len()); + // Only two values are stored, so the third step leaves the buffer. + assert!(zstd_value_offset(buf.as_slice(), 0, 3).is_err()); + Ok(()) + } + + #[test] + fn test_value_range_rejects_skipping_past_the_requested_values() { + // Frame metadata claiming more skipped values than the slice starts at would wrap. + let slice = decompressed_slice(make_interleaved(&[b"hello"]), 2, 3, 4, 1); + assert!(slice.value_range().is_err()); + assert!(slice.value_bytes().is_err()); + } + + #[rstest] + // The buffered value count matches, so both ends come from the metadata. + #[case::exact_metadata(2)] + // It does not, so the far end is walked instead. + #[case::walked_end(9)] + fn test_value_bytes_totals_the_stored_values( + #[case] n_buffered_values: usize, + ) -> VortexResult<()> { + let buf = make_interleaved(&[b"hello", b"world"]); + let slice = decompressed_slice(buf.clone(), 0, 2, 0, n_buffered_values); + let (bytes, num_bytes) = slice.value_bytes()?; + assert_eq!(bytes, buf.as_slice()); + assert_eq!(num_bytes, buf.len() - 2 * size_of::()); + Ok(()) + } + + #[test] + fn test_value_bytes_rejects_more_values_than_the_buffer_holds() { + // Frame metadata claims nine values but only two are stored. + let slice = decompressed_slice(make_interleaved(&[b"hello", b"world"]), 0, 5, 0, 9); + assert!(slice.value_bytes().is_err()); + } } diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index 607238d46b5..9198863ebb6 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #![expect(clippy::cast_possible_truncation)] +use rstest::rstest; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; @@ -11,15 +13,20 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; use vortex_array::builders::VarBinBuilder; +use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_error::VortexResult; use vortex_mask::Mask; use crate::Zstd; +use crate::ZstdArray; +use crate::ZstdData; +use crate::ZstdMetadata; #[test] fn test_zstd_compress_decompress() { @@ -236,6 +243,49 @@ fn test_zstd_append_to_offset_builder() { ); } +/// A slice decompresses whole frames, so the frames hold values on either side of the ones it +/// requests. `push_buffer_and_adjusted_views` publishes the buffers it is handed as they are, so +/// only the requested region may reach it — otherwise the finished array retains the whole frames. +#[test] +fn test_zstd_append_to_view_builder_keeps_only_the_sliced_bytes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Long enough that the values live in the data buffers rather than inline in the views. + let values = (0..12) + .map(|i| format!("value number {i} padded well past the inline limit")) + .collect::>(); + let array = VarBinViewArray::from_iter_str(&values); + // Three values per frame, so the slice below starts and ends inside a frame holding more. + let compressed = Zstd::from_var_bin_view(&array, 0, 3, &mut ctx)?.slice(4..8)?; + + // Seeded with a value of its own so the pushed buffers land after an in-progress buffer, and + // appended to twice so the second push has to rebase past the first. + let mut builder = VarBinViewBuilder::with_capacity(compressed.dtype().clone(), 9); + builder.append_value(&values[0]); + compressed.append_to_builder(&mut builder, &mut ctx)?; + compressed.append_to_builder(&mut builder, &mut ctx)?; + let appended = builder.finish_into_varbinview(); + + let expected = VarBinViewArray::from_iter_str( + std::iter::once(&values[0]) + .chain(&values[4..8]) + .chain(&values[4..8]), + ); + assert_arrays_eq!(appended, expected, &mut ctx); + + // Each stored value costs its bytes plus the u32 length prefix zstd interleaves. + let sliced_bytes: usize = values[4..8] + .iter() + .map(|value| value.len() + size_of::()) + .sum(); + let buffered: usize = appended + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().len()) + .sum(); + assert_eq!(buffered, values[0].len() + 2 * sliced_bytes); + Ok(()) +} + #[test] fn test_zstd_decompress_var_bin_view() { let mut ctx = array_session().create_execution_ctx(); @@ -277,6 +327,98 @@ fn test_sliced_array_children() { sliced.children(); } +/// Six rows, five of them stored, compressed into `values_per_frame`-sized frames with the frame +/// metadata a reader would have deserialized rewritten by `corrupt`. +fn corrupt_var_bin_view_metadata( + values_per_frame: usize, + corrupt: impl FnOnce(&mut ZstdMetadata), + ctx: &mut ExecutionCtx, +) -> VortexResult<(VarBinViewArray, ZstdArray)> { + let array = VarBinViewArray::from_iter( + [ + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + None, + Some(b"Lorem ipsum dolor sit amet".as_slice()), + Some(b"baz".as_slice()), + Some(b"quux".as_slice()), + ], + DType::Utf8(Nullability::Nullable), + ); + let mut data = ZstdData::from_var_bin_view(&array, 0, values_per_frame, ctx)?; + corrupt(&mut data.metadata); + let compressed = Zstd::try_new(array.dtype().clone(), data, array.validity()?)?; + Ok((array, compressed)) +} + +/// Frame metadata comes straight off disk, so an inconsistent value count has to surface as an +/// error from both read paths rather than a panic or a wrapped length. +#[rstest] +#[case::frame_holds_fewer_values_than_claimed(|metadata: &mut ZstdMetadata| { + metadata.frames[0].n_values = 1000; +})] +#[case::frame_value_counts_overflow(|metadata: &mut ZstdMetadata| { + metadata.frames[1].n_values = u64::MAX; +})] +#[case::missing_value_count_across_frames(|metadata: &mut ZstdMetadata| { + metadata.frames[0].n_values = 0; +})] +fn test_zstd_rejects_corrupt_frame_metadata( + #[case] corrupt: fn(&mut ZstdMetadata), +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let (_, compressed) = corrupt_var_bin_view_metadata(3, corrupt, &mut ctx)?; + + assert!(Zstd::decompress(&compressed, &mut ctx).is_err()); + + let mut builder = + VarBinBuilder::::with_capacity(compressed.dtype().clone(), compressed.len()); + assert!( + compressed + .append_to_builder(&mut builder, &mut ctx) + .is_err() + ); + Ok(()) +} + +/// Metadata written before frames recorded their value count leaves it at zero. A single frame +/// holds every stored value, so those arrays still read back. +#[test] +fn test_zstd_reads_legacy_single_frame_var_bin_metadata() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let (array, compressed) = + corrupt_var_bin_view_metadata(0, |metadata| metadata.frames[0].n_values = 0, &mut ctx)?; + + assert_arrays_eq!(compressed, array.clone().into_array(), &mut ctx); + assert_arrays_eq!( + compressed.slice(2..5)?, + array.into_array().slice(2..5)?, + &mut ctx + ); + Ok(()) +} + +/// The same legacy metadata for fixed-width values recovers the count from the frame size, which +/// stays exact across frames. +#[test] +fn test_zstd_reads_legacy_primitive_frame_metadata() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0..200_i32); + let mut data = ZstdData::from_primitive(&array, 3, 30, &mut ctx)?; + for frame in &mut data.metadata.frames { + frame.n_values = 0; + } + let compressed = Zstd::try_new(array.dtype().clone(), data, array.validity()?)?; + + assert_arrays_eq!(compressed, PrimitiveArray::from_iter(0..200_i32), &mut ctx); + assert_arrays_eq!( + compressed.slice(100..105)?, + PrimitiveArray::from_iter(100..105_i32), + &mut ctx + ); + Ok(()) +} + /// Tests that each beginning of a frame in ZSTD matches /// the buffer alignment when compressing primitive arrays. #[test] diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index 6d14cc46e86..01f8f964615 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -331,7 +331,7 @@ async fn decode_zstd(array: ZstdArray, ctx: &mut CudaExecutionCtx) -> VortexResu .indices() { AllOr::All => { - let (buffers, all_views) = reconstruct_views(&host_buffer, MAX_BUFFER_LEN); + let (buffers, all_views) = reconstruct_views(&host_buffer, 0, MAX_BUFFER_LEN); let sliced_views = all_views.slice(slice_value_idx_start..slice_value_idx_stop); Ok(Canonical::VarBinView(unsafe { From 7c9bf8775232f5bb840bb63f8a848b520c84b6dc Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 5 Aug 2026 15:58:43 +0100 Subject: [PATCH 2/5] fixes Signed-off-by: Robert Kruszewski --- encodings/zstd/src/array.rs | 116 +++++++++++++++++++++++++++++------- encodings/zstd/src/test.rs | 42 +++++++++++++ 2 files changed, 137 insertions(+), 21 deletions(-) diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 82b7b3d8a92..4562f0cac34 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -59,6 +59,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use zstd::zstd_safe::WriteBuf; @@ -331,12 +332,24 @@ where .data() .decompress_slice(array.dtype(), &unsliced_validity(array), ctx)?; let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; + append_slice_to_varbin(&slice, &mask, builder) +} + +/// Copies the values `mask` marks as valid out of `slice` and into `builder`. +fn append_slice_to_varbin( + slice: &DecompressedSlice, + mask: &Mask, + builder: &mut VarBinBuilder, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ let (values, num_bytes) = slice.value_bytes()?; // Each value is length-prefixed, so the frames can only be walked in order — which is the - // order `append_valid_slices` visits the valid rows in. A stream that runs out early yields - // empty slices and so fails the builder's byte-count check. - let mut values = zstd_values(values); - builder.append_valid_slices(num_bytes, &mask, |_| values.next().unwrap_or_default()) + // order `append_valid_slices` visits the valid rows in. A walk that runs out early hands back + // its undecoded remainder, which the builder rejects as a byte-count mismatch. + let mut values = ZstdValues::new(values); + builder.append_valid_slices(num_bytes, mask, |_| values.next_value()) } /// Hands the decompressed frames to `builder` as data buffers with views built over them. @@ -787,10 +800,16 @@ impl DecompressedSlice { Ok(start..end) } - /// The bounds within `bytes` of the length-prefixed region holding exactly this slice's values. + /// The bounds within `bytes` of the length-prefixed region that should hold exactly this + /// slice's values. /// /// Walking the length prefixes is a dependent load chain, so both ends are derived from the - /// slice metadata where possible: an unsliced array skips both walks. + /// slice metadata where possible: an unsliced array skips both walks. That makes the far end a + /// claim by the frame metadata rather than a checked fact, so a caller must still hold the + /// values it reads to the count [`Self::value_range`] gives — by walking them with + /// [`ZstdValues`], whose shortfall shows up in the byte total from [`Self::value_bytes`], or by + /// counting the ones it decodes, as [`try_reconstruct_views`] does. A region that ends part way + /// through a value is otherwise free to pass its trailing bytes off as values of their own. fn value_byte_range(&self) -> VortexResult> { let Range { start, end } = self.value_range()?; let buffer = self.bytes.as_slice(); @@ -871,22 +890,39 @@ fn zstd_value_len(buffer: &[u8], offset: usize) -> VortexResult { Ok(ViewLen::from_le_bytes(*prefix) as usize) } -/// Iterates the values of a length-prefixed region, stopping at the first prefix that leaves it. -/// -/// Stopping short leaves the value count and byte total below what the caller declared, which the -/// consuming builder rejects, so the walk does not need to report the error itself. -fn zstd_values(buffer: &[u8]) -> impl Iterator { - let mut offset = 0; - std::iter::from_fn(move || { - if offset >= buffer.len() { - return None; +/// A forward walk over the values of a length-prefixed region. +struct ZstdValues<'a> { + buffer: &'a [u8], + offset: usize, +} + +impl<'a> ZstdValues<'a> { + fn new(buffer: &'a [u8]) -> Self { + Self { buffer, offset: 0 } + } + + /// The next value, or every byte the walk could not decode once a prefix leaves the region. + /// + /// Handing back the remainder is what lets a caller that knows how many values the region holds + /// detect a walk that fell short purely from the byte total, without a second walk to validate + /// the region up front. A caller sizes that total as the region minus one length prefix per + /// value, so `k` of `n` values decoded leaves it expecting the `n - k` prefixes the walk + /// abandoned as well: the remainder either covers them and more, or is empty because the region + /// ended exactly and the decoded values already fall short of the total. Neither can add up, so + /// a shortfall is always rejected rather than passed off as trailing empty values. + fn next_value(&mut self) -> &'a [u8] { + let value_start = self.offset + size_of::(); + let value = zstd_value_len(self.buffer, self.offset) + .ok() + .and_then(|len| self.buffer.get(value_start..value_start.checked_add(len)?)); + match value { + Some(value) => { + self.offset = value_start + value.len(); + value + } + None => &self.buffer[self.offset..], } - let len = zstd_value_len(buffer, offset).ok()?; - let value_start = offset + size_of::(); - let value = buffer.get(value_start..value_start + len)?; - offset = value_start + len; - Some(value) - }) + } } impl ZstdData { @@ -1576,12 +1612,18 @@ impl OperationsVTable for Zstd { #[expect(clippy::cast_possible_truncation)] mod tests { use rstest::rstest; + use vortex_array::arrays::varbin::VarBinArrayExt as _; + use vortex_array::builders::VarBinBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; use vortex_array::validity::Validity; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; + use vortex_mask::Mask; use super::DecompressedSlice; use super::ViewLen; + use super::append_slice_to_varbin; use super::reconstruct_views; use super::try_reconstruct_views; use super::zstd_value_len; @@ -1728,4 +1770,36 @@ mod tests { let slice = decompressed_slice(make_interleaved(&[b"hello", b"world"]), 0, 5, 0, 9); assert!(slice.value_bytes().is_err()); } + + #[test] + fn test_append_to_varbin_copies_the_stored_values() -> VortexResult<()> { + let slice = decompressed_slice(make_interleaved(&[b"hello", b"world"]), 0, 2, 0, 2); + let mut builder = VarBinBuilder::::new(DType::Utf8(NonNullable)); + append_slice_to_varbin(&slice, &Mask::new_true(2), &mut builder)?; + + let appended = builder.finish_into_varbin(); + assert_eq!(appended.bytes_at(0).as_slice(), b"hello"); + assert_eq!(appended.bytes_at(1).as_slice(), b"world"); + Ok(()) + } + + /// Both ends of the region a slice reads come from the frame metadata where they can, so the + /// values in it have to be held to the value count the metadata declared. Otherwise a buffer + /// whose last value is only a length prefix reads back as a trailing empty value. + #[test] + fn test_append_to_varbin_rejects_a_dangling_length_prefix() { + let mut buffer = Vec::new(); + buffer.extend_from_slice(&3u32.to_le_bytes()); + buffer.extend_from_slice(b"cat"); + // A prefix with no value after it. It takes up exactly the four bytes the second value's + // own prefix would have, so treating that value as empty would still total the byte count + // the metadata implies and append ["cat", ""]. + buffer.extend_from_slice(&1u32.to_le_bytes()); + + let slice = decompressed_slice(ByteBuffer::copy_from(buffer.as_slice()), 0, 2, 0, 2); + let mut builder = VarBinBuilder::::new(DType::Utf8(NonNullable)); + assert!(append_slice_to_varbin(&slice, &Mask::new_true(2), &mut builder).is_err()); + // The builder rejected the values before committing any of them. + assert_eq!(builder.finish_into_varbin().len(), 0); + } } diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index 9198863ebb6..e8bded8e4a1 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -20,12 +20,14 @@ use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_mask::Mask; use crate::Zstd; use crate::ZstdArray; use crate::ZstdData; +use crate::ZstdFrameMetadata; use crate::ZstdMetadata; #[test] @@ -381,6 +383,46 @@ fn test_zstd_rejects_corrupt_frame_metadata( Ok(()) } +/// Frame bytes are as untrusted as the metadata describing them, so a frame whose last value is +/// nothing but a length prefix has to surface as an error from every read path rather than as a +/// trailing empty value. +#[test] +fn test_zstd_rejects_a_frame_ending_in_a_dangling_length_prefix() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + values.extend_from_slice(&3u32.to_le_bytes()); + values.extend_from_slice(b"cat"); + // A prefix with no value after it. It takes up exactly the four bytes the second value's own + // prefix would have, so treating that value as empty totals the byte count the metadata + // implies: two values holding three bytes between them. + values.extend_from_slice(&1u32.to_le_bytes()); + + let dtype = DType::Utf8(Nullability::NonNullable); + let compressed = Zstd::try_new( + dtype.clone(), + ZstdData::new( + None, + vec![ByteBuffer::from(zstd::bulk::compress(&values, 3)?)], + ZstdMetadata { + dictionary_size: 0, + frames: vec![ZstdFrameMetadata { + uncompressed_size: values.len() as u64, + n_values: 2, + }], + }, + 2, + ), + Validity::NonNullable, + )?; + + assert!(Zstd::decompress(&compressed, &mut ctx).is_err()); + let mut varbin = VarBinBuilder::::with_capacity(dtype.clone(), 2); + assert!(compressed.append_to_builder(&mut varbin, &mut ctx).is_err()); + let mut views = VarBinViewBuilder::with_capacity(dtype, 2); + assert!(compressed.append_to_builder(&mut views, &mut ctx).is_err()); + Ok(()) +} + /// Metadata written before frames recorded their value count leaves it at zero. A single frame /// holds every stored value, so those arrays still read back. #[test] From 4384d343e5a8bc8d4c6d8ba2978204c7fae32eb2 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 5 Aug 2026 16:00:52 +0100 Subject: [PATCH 3/5] less Signed-off-by: Robert Kruszewski --- encodings/zstd/src/array.rs | 34 ++++++++++++++-------------------- encodings/zstd/src/test.rs | 10 ++++------ 2 files changed, 18 insertions(+), 26 deletions(-) diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 4562f0cac34..4a40dd58c48 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -347,7 +347,7 @@ where let (values, num_bytes) = slice.value_bytes()?; // Each value is length-prefixed, so the frames can only be walked in order — which is the // order `append_valid_slices` visits the valid rows in. A walk that runs out early hands back - // its undecoded remainder, which the builder rejects as a byte-count mismatch. + // its remainder, which the builder rejects as a byte-count mismatch. let mut values = ZstdValues::new(values); builder.append_valid_slices(num_bytes, mask, |_| values.next_value()) } @@ -804,12 +804,11 @@ impl DecompressedSlice { /// slice's values. /// /// Walking the length prefixes is a dependent load chain, so both ends are derived from the - /// slice metadata where possible: an unsliced array skips both walks. That makes the far end a - /// claim by the frame metadata rather than a checked fact, so a caller must still hold the - /// values it reads to the count [`Self::value_range`] gives — by walking them with - /// [`ZstdValues`], whose shortfall shows up in the byte total from [`Self::value_bytes`], or by - /// counting the ones it decodes, as [`try_reconstruct_views`] does. A region that ends part way - /// through a value is otherwise free to pass its trailing bytes off as values of their own. + /// slice metadata where possible: an unsliced array skips both walks. The far end is then a + /// claim rather than a checked fact, so a caller must hold the values it reads to the count + /// from [`Self::value_range`] — with [`ZstdValues`], whose shortfall shows up in the byte total + /// from [`Self::value_bytes`], or by counting decoded values as [`try_reconstruct_views`] does. + /// Otherwise a region ending part way through a value passes its trailing bytes off as values. fn value_byte_range(&self) -> VortexResult> { let Range { start, end } = self.value_range()?; let buffer = self.bytes.as_slice(); @@ -903,13 +902,11 @@ impl<'a> ZstdValues<'a> { /// The next value, or every byte the walk could not decode once a prefix leaves the region. /// - /// Handing back the remainder is what lets a caller that knows how many values the region holds - /// detect a walk that fell short purely from the byte total, without a second walk to validate - /// the region up front. A caller sizes that total as the region minus one length prefix per - /// value, so `k` of `n` values decoded leaves it expecting the `n - k` prefixes the walk - /// abandoned as well: the remainder either covers them and more, or is empty because the region - /// ended exactly and the decoded values already fall short of the total. Neither can add up, so - /// a shortfall is always rejected rather than passed off as trailing empty values. + /// The remainder is what makes a shortfall visible in the byte total alone, without a second + /// walk to validate the region up front. A caller sizes that total as the region minus one + /// prefix per value, so `k` of `n` values decoded leaves it still expecting the `n - k` + /// prefixes the walk abandoned; the remainder covers those and more, or is empty only because + /// the region ended exactly and the decoded bytes already fall short. Neither can add up. fn next_value(&mut self) -> &'a [u8] { let value_start = self.offset + size_of::(); let value = zstd_value_len(self.buffer, self.offset) @@ -1783,17 +1780,14 @@ mod tests { Ok(()) } - /// Both ends of the region a slice reads come from the frame metadata where they can, so the - /// values in it have to be held to the value count the metadata declared. Otherwise a buffer - /// whose last value is only a length prefix reads back as a trailing empty value. #[test] fn test_append_to_varbin_rejects_a_dangling_length_prefix() { let mut buffer = Vec::new(); buffer.extend_from_slice(&3u32.to_le_bytes()); buffer.extend_from_slice(b"cat"); - // A prefix with no value after it. It takes up exactly the four bytes the second value's - // own prefix would have, so treating that value as empty would still total the byte count - // the metadata implies and append ["cat", ""]. + // A prefix with no value after it. It takes up exactly the four bytes the missing value's + // own prefix would have, so treating that value as empty still totals the byte count the + // metadata implies and appends ["cat", ""]. buffer.extend_from_slice(&1u32.to_le_bytes()); let slice = decompressed_slice(ByteBuffer::copy_from(buffer.as_slice()), 0, 2, 0, 2); diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index e8bded8e4a1..e600cff020f 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -383,18 +383,16 @@ fn test_zstd_rejects_corrupt_frame_metadata( Ok(()) } -/// Frame bytes are as untrusted as the metadata describing them, so a frame whose last value is -/// nothing but a length prefix has to surface as an error from every read path rather than as a -/// trailing empty value. +/// Frame bytes are as untrusted as the metadata describing them: a frame whose last value is +/// nothing but a length prefix has to error from every read path, not read back as an empty value. #[test] fn test_zstd_rejects_a_frame_ending_in_a_dangling_length_prefix() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let mut values = Vec::new(); values.extend_from_slice(&3u32.to_le_bytes()); values.extend_from_slice(b"cat"); - // A prefix with no value after it. It takes up exactly the four bytes the second value's own - // prefix would have, so treating that value as empty totals the byte count the metadata - // implies: two values holding three bytes between them. + // A prefix with no value after it. Treating the missing value as empty still totals the three + // bytes the metadata implies for two values. values.extend_from_slice(&1u32.to_le_bytes()); let dtype = DType::Utf8(Nullability::NonNullable); From 62fd68fe75947e565a22310bf6c9c6090481cbdb Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 5 Aug 2026 16:30:28 +0100 Subject: [PATCH 4/5] typo Signed-off-by: Robert Kruszewski --- encodings/zstd/src/array.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 4a40dd58c48..0aad8c8c4e0 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1357,7 +1357,7 @@ impl ZstdData { frame_uncompressed_size / byte_width } else { // The same fallback would read a byte count as a value count for variable-width - // values, which mis-attributes values to frames. A single frame holds every stored + // values, which misattributes values to frames. A single frame holds every stored // value, so that case is still recoverable; anything else is not. vortex_ensure!( self.frames.len() == 1, From 3d42169c3cc269bfdd105bded31271fd0da05f7f Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 5 Aug 2026 17:14:20 +0100 Subject: [PATCH 5/5] DCO Remediation Commit for Robert Kruszewski I, Robert Kruszewski , hereby add my Signed-off-by to this commit: 9fdbe75a3dd0f320873c41eb9eeebbe84eaa4a67 Signed-off-by: Robert Kruszewski