Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions encodings/fastlanes/benches/bitpack_compare.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ fn page_aligned(array: BitPackedArray) -> BitPackedArray {
parts.validity,
parts.patches,
parts.widths,
parts.chunk_offsets,
parts.len,
parts.offset,
)
Expand Down
1 change: 1 addition & 0 deletions encodings/fastlanes/benches/bitpack_compare_sweep.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ fn page_aligned(array: BitPackedArray) -> BitPackedArray {
parts.validity,
parts.patches,
parts.widths,
parts.chunk_offsets,
parts.len,
parts.offset,
)
Expand Down
10 changes: 8 additions & 2 deletions encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,12 +81,15 @@ pub fn bitpack_encode(
.transpose()?
.flatten();

let widths = ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE));
let offsets = widths.offsets_array();
let bitpacked = BitPacked::try_new(
BufferHandle::new_host(packed),
array.ptype(),
array.validity()?,
patches,
ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE)).into_array(),
widths.into_array(),
offsets,
array.len(),
0,
)?;
Expand All @@ -110,12 +113,15 @@ pub unsafe fn bitpack_encode_unchecked(
let packed = unsafe { bitpack_unchecked(&array, bit_width) };

let arr_ref = array.clone().into_array();
let widths = ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE));
let offsets = widths.offsets_array();
let bitpacked = BitPacked::try_new(
BufferHandle::new_host(packed),
array.ptype(),
array.validity()?,
None,
ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE)).into_array(),
widths.into_array(),
offsets,
array.len(),
0,
)
Expand Down
156 changes: 128 additions & 28 deletions encodings/fastlanes/src/bitpacking/array/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ pub const fn chunk_packed_bytes(bit_width: u8) -> usize {
}

/// Chunk widths and byte offsets used while encoding or executing bit-packed data.
/// Operations use this view to address each packed chunk.
/// Execution borrows the materialized children; only encoding computes prefix sums.
#[derive(Clone, Debug)]
pub struct ChunkWidths {
widths: Widths,
Expand Down Expand Up @@ -238,6 +238,10 @@ pub struct BitPackedSlots {
/// One non-nullable `u8` width per 1024-element chunk. Uniform widths use a constant array.
#[slot(4)]
pub width_table: ArrayRef,
/// Non-nullable `u64` byte boundaries, with one trailing entry after the last chunk.
/// The first offset is the origin of the packed buffer and may be nonzero after slicing.
#[slot(5)]
pub chunk_offsets: ArrayRef,
}

pub(crate) const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices {
Expand All @@ -249,6 +253,9 @@ pub(crate) const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices {
/// The dtype of the width table child: one byte per chunk.
pub(crate) const WIDTH_TABLE_DTYPE: DType = DType::Primitive(PType::U8, Nullability::NonNullable);

pub(crate) const CHUNK_OFFSETS_DTYPE: DType =
DType::Primitive(PType::U64, Nullability::NonNullable);

impl IntoArray for ChunkWidths {
fn into_array(self) -> ArrayRef {
if self.is_uniform() {
Expand All @@ -259,23 +266,34 @@ impl IntoArray for ChunkWidths {
}
}

/// Read the width child without executing it during reduction.
pub(crate) fn materialized_widths(table: &ArrayRef) -> VortexResult<Option<ChunkWidths>> {
if let Some(constant) = table.as_opt::<Constant>() {
return Ok(Some(ChunkWidths::uniform(
u8::try_from(constant.scalar())?,
table.len(),
)));
}
Ok(table
/// Read materialized children without executing them during reduction.
pub(crate) fn materialized_widths(
table: &ArrayRef,
offsets: &ArrayRef,
) -> VortexResult<Option<ChunkWidths>> {
let widths = if let Some(constant) = table.as_opt::<Constant>() {
Widths::Uniform {
width: u8::try_from(constant.scalar())?,
len: table.len(),
}
} else if let Some(primitive) = table
.as_opt::<Primitive>()
.filter(|a| a.buffer_handle().is_on_host())
{
Widths::PerChunk(primitive.to_buffer::<u8>())
} else {
return Ok(None);
};
Ok(offsets
.as_opt::<Primitive>()
.filter(|a| a.buffer_handle().is_on_host())
.map(|a| ChunkWidths::new(a.to_buffer::<u8>())))
.map(|a| ChunkWidths::from_buffers(widths, a.to_buffer::<u64>())))
}

pub struct BitPackedDataParts {
pub offset: u16,
pub widths: ArrayRef,
pub chunk_offsets: ArrayRef,
pub len: usize,
pub packed: BufferHandle,
pub patches: Option<Patches>,
Expand Down Expand Up @@ -324,9 +342,11 @@ impl BitPackedData {
/// * `validity` must have `length` len
/// * Any patches must have any `array_len` equal to `length`
/// * The width-table child must hold one non-nullable `u8` per chunk.
/// * The offsets child must hold `num_chunks + 1` non-nullable `u64` byte boundaries.
///
/// Once the widths are materialized, they must be no wider than `ptype`, and the packed
/// buffer must be exactly the sum of the chunks' packed sizes. Compressed children are checked at execution time, before unpacking.
/// buffer must be exactly the sum of the chunks' packed sizes. Offset differences must
/// match the widths. Compressed children are checked at execution time, before unpacking.
///
/// Any violation of these preconditions will result in an error.
pub fn try_new(
Expand All @@ -352,6 +372,7 @@ impl BitPackedData {
validity: &Validity,
patches: Option<&Patches>,
table: &ArrayRef,
offsets: &ArrayRef,
length: usize,
) -> VortexResult<()> {
vortex_ensure!(ptype.is_int(), MismatchedTypes: "integer", ptype);
Expand All @@ -378,8 +399,19 @@ impl BitPackedData {
"Expected {num_chunks} chunk widths, got {}",
table.len()
);
vortex_ensure!(
offsets.dtype() == &CHUNK_OFFSETS_DTYPE,
"BitPacked chunk offsets must be {CHUNK_OFFSETS_DTYPE}, got {}",
offsets.dtype()
);
vortex_ensure!(
offsets.len() == num_chunks + 1,
"Expected {} chunk offsets, got {}",
num_chunks + 1,
offsets.len()
);
// Compressed children are checked once materialized, before any unchecked unpacking.
if let Some(widths) = materialized_widths(table)? {
if let Some(widths) = materialized_widths(table, offsets)? {
Self::validate_widths(&self.packed, ptype, &widths)?;
}
Ok(())
Expand Down Expand Up @@ -504,24 +536,40 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt {
BitPackedData::packed(self)
}

/// Prepare and validate the width child once for a bulk operation.
/// Prepare and validate both children once for a bulk operation, without computing prefix sums.
fn chunk_widths(&self, ctx: &mut ExecutionCtx) -> VortexResult<ChunkWidths> {
let widths = match materialized_widths(self.width_table())? {
let widths = match materialized_widths(self.width_table(), self.chunk_offsets())? {
Some(widths) => widths,
None => ChunkWidths::new(
self.width_table()
None => {
let table = self.width_table();
let widths = if let Some(constant) = table.as_opt::<Constant>() {
Widths::Uniform {
width: u8::try_from(constant.scalar())?,
len: table.len(),
}
} else {
Widths::PerChunk(
table
.clone()
.execute::<PrimitiveArray>(ctx)?
.to_buffer::<u8>(),
)
};
let offsets = self
.chunk_offsets()
.clone()
.execute::<PrimitiveArray>(ctx)?
.to_buffer::<u8>(),
),
.to_buffer::<u64>();
ChunkWidths::from_buffers(widths, offsets)
}
};
BitPackedData::validate_widths(self.packed(), self.as_ref().dtype().as_ptype(), &widths)?;
Ok(widths)
}

/// Read and validate widths only when the child is already materialized.
/// Read and validate widths and offsets only when their children are already materialized.
fn materialized_chunk_widths(&self) -> VortexResult<Option<ChunkWidths>> {
let widths = materialized_widths(self.width_table())?;
let widths = materialized_widths(self.width_table(), self.chunk_offsets())?;
if let Some(widths) = &widths {
BitPackedData::validate_widths(
self.packed(),
Expand All @@ -532,18 +580,70 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt {
Ok(widths)
}

/// Locate and validate one chunk using the width table.
/// Read one byte boundary without executing the entire offsets child.
fn chunk_byte_offset(&self, boundary: usize, ctx: &mut ExecutionCtx) -> VortexResult<u64> {
vortex_ensure!(
boundary < self.chunk_offsets().len(),
"Chunk boundary out of bounds"
);
if let Some(offsets) = self
.chunk_offsets()
.as_opt::<Primitive>()
.filter(|a| a.buffer_handle().is_on_host())
{
Ok(offsets.as_slice::<u64>()[boundary])
} else {
u64::try_from(&self.chunk_offsets().execute_scalar(boundary, ctx)?)
}
}

/// Locate and validate one chunk using scalar child access, without materializing the tables.
fn chunk_range(
&self,
chunk: usize,
ctx: &mut ExecutionCtx,
) -> VortexResult<(Range<usize>, u8)> {
let widths = self.chunk_widths(ctx)?;
vortex_ensure!(chunk < widths.len(), "Chunk index out of bounds");
Ok((
widths.byte_offset(chunk)..widths.byte_offset(chunk + 1),
widths.width(chunk),
))
vortex_ensure!(
chunk < self.width_table().len(),
"Chunk index out of bounds"
);
let width = if let Some(table) = self
.width_table()
.as_opt::<Primitive>()
.filter(|a| a.buffer_handle().is_on_host())
{
table.as_slice::<u8>()[chunk]
} else if let Some(table) = self.width_table().as_opt::<Constant>() {
u8::try_from(table.scalar())?
} else {
u8::try_from(&self.width_table().execute_scalar(chunk, ctx)?)?
};
vortex_ensure!(
width as usize <= self.as_ref().dtype().as_ptype().bit_width(),
"Unsupported bit width {width}"
);
let base = self.chunk_byte_offset(0, ctx)?;
let start = if chunk == 0 {
base
} else {
self.chunk_byte_offset(chunk, ctx)?
};
let end = self.chunk_byte_offset(chunk + 1, ctx)?;
vortex_ensure!(
end.checked_sub(start) == Some(chunk_packed_bytes(width) as u64),
"Chunk {chunk} offsets do not match its bit width"
);
let start = start
.checked_sub(base)
.ok_or_else(|| vortex_err!("Chunk offset precedes buffer origin"))?;
let end = end
.checked_sub(base)
.ok_or_else(|| vortex_err!("Chunk offset precedes buffer origin"))?;
vortex_ensure!(
start % (FL_CHUNK_SIZE / 8) as u64 == 0 && end <= self.packed().len() as u64,
"Chunk offsets are unaligned or exceed the packed buffer"
);
Ok((usize::try_from(start)?..usize::try_from(end)?, width))
}

#[inline]
Expand Down
53 changes: 53 additions & 0 deletions encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,16 @@

use std::sync::LazyLock;

use rstest::rstest;
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::arrays::Constant;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::slice::SliceKernel;
use vortex_array::assert_arrays_eq;
use vortex_array::scalar::Scalar;
use vortex_buffer::Buffer;
use vortex_buffer::buffer;
use vortex_error::VortexResult;
use vortex_session::VortexSession;

Expand Down Expand Up @@ -51,6 +56,54 @@ fn encode(values: &[u32]) -> VortexResult<BitPackedArray> {
bitpack_to_best_bit_width(&PrimitiveArray::from_iter(values.iter().copied()), &mut ctx)
}

#[rstest]
#[case::decreasing(buffer![0u64, 128, 256, 128])]
#[case::wrong_width(buffer![0u64, 128, 256, 385])]
#[case::out_of_bounds(buffer![0u64, 128, 256, u64::MAX])]
fn invalid_offsets_rejected_before_unpacking(#[case] offsets: Buffer<u64>) -> VortexResult<()> {
let mut ctx = SESSION.create_execution_ctx();
let values = PrimitiveArray::from_iter((0..3072u32).map(|i| i % 2));
let packed = bitpack_to_best_bit_width(&values, &mut ctx)?;
let widths = packed.width_table().clone();
let offsets = offsets.into_array();
assert!(BitPacked::with_chunk_layout(packed.clone(), widths.clone(), offsets.clone()).is_err());
let offsets = offsets.execute::<PrimitiveArray>(&mut ctx)?;
let offsets = bitpack_to_best_bit_width(&offsets, &mut ctx)?.into_array();
let packed = BitPacked::with_chunk_layout(packed, widths, offsets)?.into_array();
// An isolated scalar checks only its own chunk; bulk unpacking validates the whole layout.
assert_eq!(packed.execute_scalar(1, &mut ctx)?, Scalar::from(1u32));
assert!(packed.execute_scalar(2048, &mut ctx).is_err());
assert!(packed.execute::<PrimitiveArray>(&mut ctx).is_err());
Ok(())
}

#[test]
fn slice_rejects_unaligned_offsets() -> VortexResult<()> {
let mut ctx = SESSION.create_execution_ctx();
let values = PrimitiveArray::from_iter((0..3072u32).map(|i| i % 2));
let packed = bitpack_to_best_bit_width(&values, &mut ctx)?;
let widths = packed.width_table().clone();
let offsets = PrimitiveArray::from_iter([0u64, 127, 255, 383]);
let offsets = bitpack_to_best_bit_width(&offsets, &mut ctx)?.into_array();
let packed = BitPacked::with_chunk_layout(packed, widths, offsets)?;
assert!(<BitPacked as SliceKernel>::slice(packed.as_view(), 1024..2048, &mut ctx).is_err());
Ok(())
}

#[test]
fn offset_child_shape_is_validated() -> VortexResult<()> {
let packed = encode(&varied(100))?;
let widths = packed.width_table().clone();
assert!(
BitPacked::with_chunk_layout(packed.clone(), widths.clone(), buffer![0u64].into_array())
.is_err()
);
let wrong_dtype =
PrimitiveArray::from_iter(vec![0u32; packed.chunk_offsets().len()]).into_array();
assert!(BitPacked::with_chunk_layout(packed, widths, wrong_dtype).is_err());
Ok(())
}

/// Every array carries a non-nullable `u8` width per chunk, including uniform arrays.
#[test]
fn width_table_is_validated() -> VortexResult<()> {
Expand Down
1 change: 1 addition & 0 deletions encodings/fastlanes/src/bitpacking/compute/cast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ fn build_with_validity(
.map(|patches| patches.map_values(|values| values.cast(dtype.clone())))
.transpose()?,
array.width_table().clone(),
array.chunk_offsets().clone(),
array.len(),
array.offset(),
)?
Expand Down
Loading
Loading