From 377f70c71cf0b8cb98c5cd9861c1f90688ca2c73 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 18 Sep 2026 15:19:05 -0400 Subject: [PATCH] feat: add bitpacked v2 serialization Signed-off-by: "Matt Katz" --- .../src/bitpacking/chunk_widths_tests.rs | 332 +++++++++++++++++- encodings/fastlanes/src/bitpacking/mod.rs | 4 +- .../src/bitpacking/plugin/bitpacked.rs | 294 +++++++++++----- .../fastlanes/src/bitpacking/plugin/mod.rs | 10 +- .../src/bitpacking/plugin/patched.rs | 12 +- .../fastlanes/src/bitpacking/serde_tests.rs | 83 ----- 6 files changed, 550 insertions(+), 185 deletions(-) delete mode 100644 encodings/fastlanes/src/bitpacking/serde_tests.rs diff --git a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs index eeb00917c3e..1600a6491bb 100644 --- a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs +++ b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs @@ -5,7 +5,14 @@ use std::sync::LazyLock; +use prost::Message; use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Constant; @@ -14,21 +21,40 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SliceArray; use vortex_array::arrays::slice::SliceKernel; use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; +use vortex_array::compute::conformance::cast::test_cast_conformance; +use vortex_array::compute::conformance::consistency::test_array_consistency; +use vortex_array::compute::conformance::filter::test_filter_conformance; +use vortex_array::compute::conformance::take::test_take_conformance; +use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; +use vortex_array::serde::ArrayChildren; +use vortex_array::serde::SerializeOptions; +use vortex_array::serde::SerializedArray; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; use crate::BitPacked; use crate::BitPackedArray; use crate::BitPackedArrayExt; use crate::BitPackedArraySlotsExt; +use crate::BitPackedPlugin; use crate::ChunkWidths; use crate::FL_CHUNK_SIZE; +use crate::bitpacked_v2_id; use crate::bitpacking::bitpack_compress::bitpack_encode_with_widths; use crate::bitpacking::bitpack_compress::bitpack_to_best_bit_width; +use crate::bitpacking::plugin::BitPackedMetadata; +use crate::bitpacking::plugin::BitPackedV2Metadata; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -74,6 +100,10 @@ fn encode(values: &[u32]) -> VortexResult { ) } +fn primitive(values: &[u32]) -> ArrayRef { + PrimitiveArray::from_iter(values.iter().copied()).into_array() +} + #[test] fn explicit_widths_including_full_width_chunk() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -101,6 +131,226 @@ fn explicit_widths_including_full_width_chunk() -> VortexResult<()> { Ok(()) } +/// Serialize `array` through the session and read it back through the plugin registered for the +/// serialized ID, as a file reader would. +fn serde_roundtrip(array: &BitPackedArray) -> VortexResult<(ArrayId, Vec, ArrayRef)> { + let array_ref = array.as_array(); + let serialization = SESSION + .array_serialize(array_ref)? + .ok_or_else(|| vortex_err!("BitPacked must serialize"))?; + let array_ctx = ArrayContext::empty(); + let buffers = array_ref.serialize(&array_ctx, &SESSION, &SerializeOptions::default())?; + let mut bytes = ByteBufferMut::empty(); + for buffer in buffers { + bytes.extend_from_slice(&buffer); + } + let read = SerializedArray::try_from(bytes.freeze())?.decode( + array_ref.dtype(), + array_ref.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + )?; + Ok((serialization.serialized_id, serialization.metadata, read)) +} + +/// Differing chunk widths serialize under the v2 ID with the width table as a child, and read +/// back with the same widths. +#[test] +fn differing_widths_serialize_as_v2() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?; + let table = packed.width_table(); + assert_eq!( + table + .clone() + .execute::(&mut ctx)? + .as_slice::(), + packed + .chunk_widths(&mut SESSION.create_execution_ctx())? + .as_buffer() + .as_slice() + ); + let (id, metadata, read) = serde_roundtrip(&packed)?; + assert_eq!(id, bitpacked_v2_id()); + assert_eq!(BitPackedV2Metadata::decode(metadata.as_slice())?.offset, 0); + assert_eq!( + read.as_::() + .chunk_widths(&mut SESSION.create_execution_ctx())?, + packed.chunk_widths(&mut SESSION.create_execution_ctx())? + ); + assert_eq!( + read.as_::().width_table().len(), + packed.width_table().len() + ); + assert_arrays_eq!(read, primitive(&values), &mut ctx); + Ok(()) +} + +/// One shared width serializes under the original ID with the original metadata and no width +/// table, byte for byte. +#[test] +fn uniform_widths_serialize_as_original_format() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..3000).map(|i| i % 128).collect(); + let packed = encode(&values)?; + assert!(packed.width_table().is::()); + assert_eq!(packed.as_array().children().len(), 2); + assert!( + SESSION + .array_serialize(packed.as_array())? + .ok_or_else(|| vortex_err!("must serialize"))? + .children + .is_empty() + ); + let (id, metadata, read) = serde_roundtrip(&packed)?; + assert_eq!(id, ArrayVTable::id(&BitPacked)); + let original = BitPackedMetadata { + bit_width: 7, + offset: 0, + patches: None, + } + .encode_to_vec(); + assert_eq!(metadata, original); + assert_arrays_eq!(read, primitive(&values), &mut ctx); + Ok(()) +} + +/// An array with no chunks has nothing to tabulate and stays in the original format. +#[test] +fn empty_array_serializes_as_original_format() -> VortexResult<()> { + let packed = encode(&[])?; + assert!(packed.width_table().is::()); + let (id, _, read) = serde_roundtrip(&packed)?; + assert_eq!(id, ArrayVTable::id(&BitPacked)); + assert!(read.is_empty()); + Ok(()) +} + +/// A compressor may re-encode the width table. The re-encoded child survives a round trip and +/// still yields the same widths. +#[test] +fn compressed_width_table_round_trips() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?; + let table = PrimitiveArray::new( + packed + .chunk_widths(&mut SESSION.create_execution_ctx())? + .as_buffer(), + Validity::NonNullable, + ); + let compressed_table = bitpack_to_best_bit_width(&table, &mut ctx)?.into_array(); + assert!(compressed_table.is::()); + let packed = BitPacked::with_width_table(packed, compressed_table)?; + let (id, _, read) = serde_roundtrip(&packed)?; + assert_eq!(id, bitpacked_v2_id()); + let view = read.as_::(); + assert_eq!( + view.chunk_widths(&mut SESSION.create_execution_ctx())?, + packed.chunk_widths(&mut SESSION.create_execution_ctx())? + ); + assert!(view.width_table().is::()); + assert_arrays_eq!(read, primitive(&values), &mut ctx); + Ok(()) +} + +#[test] +fn widths_must_agree_with_offsets() -> VortexResult<()> { + let packed = encode(&varied(0))?; + let table = PrimitiveArray::from_iter( + packed + .chunk_widths(&mut SESSION.create_execution_ctx())? + .as_buffer() + .as_slice() + .iter() + .rev() + .copied(), + ) + .into_array(); + assert!(BitPacked::with_width_table(packed, table).is_err()); + Ok(()) +} + +#[rstest] +#[case::too_wide(64u8)] +#[case::wrong_packed_size(1u8)] +fn compressed_widths_are_validated_on_execution(#[case] width: u8) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let packed = encode(&varied(0))?; + let mut serialized = SESSION + .array_serialize(packed.as_array())? + .ok_or_else(|| vortex_err!("must serialize"))?; + let table = PrimitiveArray::from_iter(vec![width; packed.width_table().len()]); + let table = bitpack_to_best_bit_width(&table, &mut ctx)?.into_array(); + let table_idx = serialized.children.len() - 2; + serialized.children[table_idx] = table; + let buffers: Vec<_> = serialized + .buffers + .into_iter() + .map(BufferHandle::new_host) + .collect(); + let read = BitPackedPlugin.deserialize( + ArrayDeserialization::new( + serialized.serialized_id, + packed.dtype(), + packed.len(), + &serialized.metadata, + &buffers, + &StrictChildren(serialized.children), + ), + &SESSION, + )?; + assert!(read.as_::().width_table().is::()); + assert!(read.execute::(&mut ctx).is_err()); + Ok(()) +} + +#[rstest] +#[case::widths(true, false)] +#[case::offsets(false, true)] +#[case::both(true, true)] +fn compressed_layout_kernels( + #[case] compress_widths: bool, + #[case] compress_offsets: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let packed = encode(&varied(100))?; + let compress = |child: &ArrayRef, enabled: bool| -> VortexResult { + if !enabled { + return Ok(child.clone()); + } + let mut ctx = SESSION.create_execution_ctx(); + let primitive = child.clone().execute::(&mut ctx)?; + Ok(bitpack_to_best_bit_width(&primitive, &mut ctx)?.into_array()) + }; + let widths = compress(packed.width_table(), compress_widths)?; + let offsets = compress(packed.chunk_offsets(), compress_offsets)?; + let packed = BitPacked::with_chunk_layout(packed, widths, offsets)?; + let (_, _, read) = serde_roundtrip(&packed)?; + assert_eq!( + read.as_::().width_table().is::(), + compress_widths + ); + assert_eq!( + read.as_::().chunk_offsets().is::(), + compress_offsets + ); + assert_arrays_eq!(read, primitive(&varied(100)), &mut ctx); + let packed = packed.into_array(); + test_array_consistency(&packed, &mut ctx); + test_take_conformance(&packed, &mut ctx); + test_filter_conformance(&packed, &mut ctx); + test_cast_conformance(&packed, &mut ctx); + test_binary_numeric_array(&packed, &mut ctx); + assert_arrays_eq!( + packed.slice(900..2100)?, + primitive(&varied(100)).slice(900..2100)?, + &mut ctx + ); + Ok(()) +} + #[test] fn slice_preserves_offset_origin() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -142,7 +392,7 @@ fn slice_preserves_offset_origin() -> VortexResult<()> { .as_slice::()[1..] .as_ptr() ); - let read = sliced.clone(); + let (_, _, read) = serde_roundtrip(&bp.into_owned())?; assert_arrays_eq!( read, values.clone().into_array().slice(1100..2300)?, @@ -226,3 +476,83 @@ fn width_table_is_validated() -> VortexResult<()> { assert_arrays_eq!(uniform, replaced, &mut SESSION.create_execution_ctx()); Ok(()) } + +/// Children that report a dtype or length mismatch as an error, as a file reader does, instead +/// of panicking like the slice implementation. +struct StrictChildren(Vec); + +impl ArrayChildren for StrictChildren { + fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult { + let child = + <[ArrayRef]>::get(&self.0, index).ok_or_else(|| vortex_err!("no child {index}"))?; + vortex_ensure!( + child.dtype() == dtype, + "child {index} has dtype {}, expected {dtype}", + child.dtype() + ); + vortex_ensure!( + child.len() == len, + "child {index} has length {}, expected {len}", + child.len() + ); + Ok(child.clone()) + } + + fn len(&self) -> usize { + self.0.len() + } +} + +/// The plugin owns wire children, which differ from in-memory slots for the v1 format. +#[test] +fn bare_vtable_requires_plugin() -> VortexResult<()> { + let uniform = encode(&(0..3000u32).map(|i| i % 128).collect::>())?; + assert!(ArrayVTable::serialize(uniform.as_view(), &SESSION).is_err()); + let differing = encode(&varied(100))?; + assert!(ArrayVTable::serialize(differing.as_view(), &SESSION).is_err()); + Ok(()) +} + +/// Each ID keeps its contract: the original ID cannot read children that carry a width table, +/// and the v2 ID demands one. +#[test] +fn each_format_keeps_its_contract() -> VortexResult<()> { + let read_as = |array: &BitPackedArray, id: ArrayId| -> VortexResult<()> { + let array_ref = array.as_array(); + let serialization = SESSION + .array_serialize(array_ref)? + .ok_or_else(|| vortex_err!("BitPacked must serialize"))?; + let children = StrictChildren(serialization.children.clone()); + let buffers = serialization + .buffers + .clone() + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + ArrayPlugin::deserialize( + &BitPackedPlugin, + ArrayDeserialization::new( + id, + array_ref.dtype(), + array_ref.len(), + &serialization.metadata, + &buffers, + &children, + ), + &SESSION, + ) + .map(|_| ()) + }; + + let differing = encode(&varied(100))?; + assert!( + read_as(&differing, ArrayVTable::id(&BitPacked)).is_err(), + "the original ID must reject a width table" + ); + let uniform = encode(&(0..3000u32).map(|i| i % 128).collect::>())?; + assert!( + read_as(&uniform, bitpacked_v2_id()).is_err(), + "the v2 ID must demand a width table" + ); + Ok(()) +} diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index 8ce20e5f160..24864c0c3e1 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -20,6 +20,7 @@ mod vtable; pub(crate) use plugin::BitPackedPatchedPlugin; pub use plugin::BitPackedPlugin; +pub use plugin::bitpacked_v2_id; pub use vtable::BitPacked; pub use vtable::BitPackedArray; @@ -27,8 +28,5 @@ pub(crate) fn initialize(session: &vortex_session::VortexSession) { vtable::initialize(session); } -#[cfg(test)] -mod serde_tests; - #[cfg(test)] mod chunk_widths_tests; diff --git a/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs b/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs index 2d1fe66417e..b18500d0774 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Serialization plugin for the frozen bit-packed wire format. +//! Bit-packed wire formats and their serialization plugin. use prost::Message; use vortex_array::Array; @@ -15,9 +15,12 @@ use vortex_array::ArraySlots; use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; use vortex_array::patches::Patches; use vortex_array::patches::PatchesData; use vortex_array::patches::PatchesMetadata; +use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::validity_to_child; use vortex_error::VortexResult; @@ -25,6 +28,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::BitPacked; use crate::BitPackedArrayExt; @@ -32,6 +36,8 @@ use crate::BitPackedData; use crate::ChunkWidths; use crate::FL_CHUNK_SIZE; use crate::bitpacking::array::BitPackedSlots; +use crate::bitpacking::array::CHUNK_OFFSETS_DTYPE; +use crate::bitpacking::array::WIDTH_TABLE_DTYPE; /// Metadata of the frozen `fastlanes.bitpacked` wire format. #[derive(Clone, prost::Message)] @@ -44,7 +50,139 @@ pub(crate) struct BitPackedMetadata { pub(crate) patches: Option, } -/// Serialization boundary for the frozen `fastlanes.bitpacked` wire format. +/// Read the frozen v1 metadata, packed buffer, and patch/validity children. +fn deserialize_v1(parts: ArrayDeserialization<'_>) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + + let metadata = BitPackedMetadata::decode(metadata)?; + let packed = single_buffer(buffers)?; + let (patches, validity, _) = deserialize_children(children, metadata.patches, dtype, len, 0)?; + let bit_width = u8::try_from(metadata.bit_width).map_err(|_| { + vortex_err!( + "BitPackedMetadata bit_width {} does not fit in u8", + metadata.bit_width + ) + })?; + let offset = offset_from_metadata(metadata.offset)?; + let num_chunks = (len + offset as usize).div_ceil(FL_CHUNK_SIZE); + + let slots = { + let mut s = ArraySlots::with_capacity(BitPackedSlots::COUNT); + PatchesData::push_slots(&mut s, patches.as_ref()); + s.push(validity_to_child(&validity, len)); + let widths = ChunkWidths::uniform(bit_width, num_chunks); + let offsets = widths.offsets_array(); + s.push(Some(widths.into_array())); + s.push(Some(offsets)); + s + }; + let data = BitPackedData::try_new(packed, patches, offset)?; + Ok(Array::::try_from_parts( + ArrayParts::new(BitPacked, dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) +} + +/// The single packed buffer of a serialized bit-packed array. +fn single_buffer(buffers: &[BufferHandle]) -> VortexResult { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + Ok(buffers[0].clone()) +} + +/// The offset into the first chunk, which the metadata stores as a `u32`. +fn offset_from_metadata(offset: u32) -> VortexResult { + u16::try_from(offset) + .map_err(|_| vortex_err!("BitPackedMetadata offset {offset} does not fit in u16")) +} + +/// Read the patches and validity children that both wire formats share. +/// +/// Children run: the patches, then a validity bitmap if there is one, then `trailing` children +/// the caller reads itself. Returns the index of the first trailing child. +fn deserialize_children( + children: &dyn ArrayChildren, + patches: Option, + dtype: &DType, + len: usize, + trailing: usize, +) -> VortexResult<(Option, Validity, usize)> { + let num_patch_children = match &patches { + None => 0, + Some(patches_meta) if patches_meta.chunk_offsets_dtype()?.is_some() => 3, + Some(_) => 2, + }; + let num_fixed = num_patch_children + trailing; + let has_validity = match children.len().checked_sub(num_fixed) { + Some(0) => false, + Some(1) => true, + _ => vortex_bail!( + "Expected {num_fixed} or {} children, got {}", + num_fixed + 1, + children.len() + ), + }; + let validity = if has_validity { + Validity::Array(children.get(num_patch_children, &Validity::DTYPE, len)?) + } else { + Validity::from(dtype.nullability()) + }; + let patches = patches + .map(|p| { + let indices = children.get(0, &p.indices_dtype()?, p.len()?)?; + let values = children.get(1, dtype, p.len()?)?; + let chunk_offsets = p + .chunk_offsets_dtype()? + .map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize)) + .transpose()?; + Patches::new(len, p.offset()?, indices, values, chunk_offsets) + }) + .transpose()?; + Ok(( + patches, + validity, + num_patch_children + usize::from(has_validity), + )) +} + +/// The serialized format for arrays whose chunks do not all share one bit width. +/// +/// The original `fastlanes.bitpacked` format carries a single `bit_width`, and readers of that +/// format assume every chunk uses it. Arrays with differing chunk widths therefore serialize under +/// this successor ID, which older readers reject as unknown instead of misreading. +pub fn bitpacked_v2_id() -> ArrayId { + static ID: CachedId = CachedId::new("fastlanes.bitpacked_v2"); + *ID +} + +/// Metadata of the `fastlanes.bitpacked_v2` format. Chunk widths and byte offsets travel +/// in separate children, keeping the metadata bounded. +/// +/// Tag 1 is left unused: it is `bit_width` in the original format, so metadata misdirected across +/// the two IDs decodes to the right fields and fails on the child layout instead. +#[derive(Clone, prost::Message)] +pub(crate) struct BitPackedV2Metadata { + #[prost(uint32, tag = "2")] + pub(crate) offset: u32, + #[prost(message, optional, tag = "3")] + pub(crate) patches: Option, +} + +/// The [`ArrayPlugin`] for `BitPacked`, owning both of its wire formats. +/// +/// Arrays whose chunks share one width serialize without either layout child as the +/// frozen `fastlanes.bitpacked` format, byte for byte. Differing widths serialize as +/// `fastlanes.bitpacked_v2`, with the width table followed by the byte-offset boundaries. #[derive(Debug, Clone)] pub struct BitPackedPlugin; @@ -53,6 +191,10 @@ impl ArrayPlugin for BitPackedPlugin { ArrayVTable::id(&BitPacked) } + fn serialized_ids(&self) -> Vec { + vec![self.id(), bitpacked_v2_id()] + } + fn serialize( &self, array: &ArrayRef, @@ -66,11 +208,7 @@ impl ArrayPlugin for BitPackedPlugin { ); let view = array.as_::(); let widths = view.chunk_widths(&mut session.create_execution_ctx())?; - vortex_ensure!( - widths.is_uniform(), - "Nonuniform widths require the v2 wire format" - ); - { + if widths.is_uniform() { let metadata = BitPackedMetadata { bit_width: widths.max_width() as u32, offset: view.offset() as u32, @@ -85,110 +223,82 @@ impl ArrayPlugin for BitPackedPlugin { .flatten() .cloned() .collect(); - Ok(Some(ArraySerialization::new( + return Ok(Some(ArraySerialization::new( self.id(), metadata, array.buffers(), children, - ))) + ))); + } + let metadata = BitPackedV2Metadata { + offset: view.offset() as u32, + patches: view + .patches() + .map(|p| p.to_metadata(view.len(), view.dtype())) + .transpose()?, } + .encode_to_vec(); + // The children run patches, validity, width table, then chunk offsets. + Ok(Some(ArraySerialization::from_array( + bitpacked_v2_id(), + array, + metadata, + ))) } fn deserialize( &self, parts: ArrayDeserialization<'_>, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult { + if parts.serialized_id == self.id() { + return deserialize_v1(parts); + } vortex_ensure!( - self.id() == parts.serialized_id, + parts.serialized_id == bitpacked_v2_id(), "array plugin {} does not recognize serialized ID {}", self.id(), parts.serialized_id, ); - let ArrayDeserialization { - dtype, - len, - metadata, - buffers, - children, - .. - } = parts; - - let metadata = BitPackedMetadata::decode(metadata)?; - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let packed = buffers[0].clone(); - - let load_validity = |child_idx: usize| { - if children.len() == child_idx { - Ok(Validity::from(dtype.nullability())) - } else if children.len() == child_idx + 1 { - let validity = children.get(child_idx, &Validity::DTYPE, len)?; - Ok(Validity::Array(validity)) - } else { - vortex_bail!( - "Expected {} or {} children, got {}", - child_idx, - child_idx + 1, - children.len() - ); - } - }; - - let validity_idx = match &metadata.patches { - None => 0, - Some(patches_meta) if patches_meta.chunk_offsets_dtype()?.is_some() => 3, - Some(_) => 2, - }; - - let validity = load_validity(validity_idx)?; - - let patches = metadata - .patches - .map(|p| { - let indices = children.get(0, &p.indices_dtype()?, p.len()?)?; - let values = children.get(1, dtype, p.len()?)?; - let chunk_offsets = p - .chunk_offsets_dtype()? - .map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize)) - .transpose()?; - - Patches::new(len, p.offset()?, indices, values, chunk_offsets) - }) - .transpose()?; - - let bit_width = u8::try_from(metadata.bit_width).map_err(|_| { - vortex_err!( - "BitPackedMetadata bit_width {} does not fit in u8", - metadata.bit_width - ) - })?; - let offset = u16::try_from(metadata.offset).map_err(|_| { - vortex_err!( - "BitPackedMetadata offset {} does not fit in u16", - metadata.offset - ) - })?; - let num_chunks = (len + offset as usize).div_ceil(FL_CHUNK_SIZE); - let slots = { - let mut s = ArraySlots::with_capacity(BitPackedSlots::COUNT); - PatchesData::push_slots(&mut s, patches.as_ref()); - s.push(validity_to_child(&validity, len)); - let widths = ChunkWidths::uniform(bit_width, num_chunks); - let offsets = widths.offsets_array(); - s.push(Some(widths.into_array())); - s.push(Some(offsets)); - s - }; - let data = BitPackedData::try_new(packed, patches, offset)?; - Ok(Array::::try_from_parts( - ArrayParts::new(BitPacked, dtype.clone(), len, data).with_slots(slots), - )? - .into_array()) + deserialize_v2(parts, session) } } +/// Read the `fastlanes.bitpacked_v2` format: [`BitPackedV2Metadata`], one packed buffer, and +/// children running patches, validity, width table, then chunk offsets. +fn deserialize_v2( + parts: ArrayDeserialization<'_>, + _session: &VortexSession, +) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = BitPackedV2Metadata::decode(metadata)?; + let packed = single_buffer(buffers)?; + let offset = offset_from_metadata(metadata.offset)?; + let num_chunks = (len + offset as usize).div_ceil(FL_CHUNK_SIZE); + let (patches, validity, table_idx) = + deserialize_children(children, metadata.patches, dtype, len, 2)?; + let table = children.get(table_idx, &WIDTH_TABLE_DTYPE, num_chunks)?; + let offsets = children.get(table_idx + 1, &CHUNK_OFFSETS_DTYPE, num_chunks + 1)?; + Ok(BitPacked::try_new( + packed, + dtype.as_ptype(), + validity, + patches, + table, + offsets, + len, + offset, + )? + .into_array()) +} + #[cfg(test)] mod tests { use prost::Message; diff --git a/encodings/fastlanes/src/bitpacking/plugin/mod.rs b/encodings/fastlanes/src/bitpacking/plugin/mod.rs index d97594f7c26..55d933f7785 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/mod.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/mod.rs @@ -1,7 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Serialization plugins for the frozen bit-packed format and its external-patches adapter. +//! [`vortex_array::ArrayPlugin`]s for bit-packed arrays. +//! +//! [`BitPackedPlugin`] owns the wire history of `BitPacked`: the frozen `fastlanes.bitpacked` +//! format for arrays whose chunks share one width, and `fastlanes.bitpacked_v2`, whose width +//! table and offsets children describe each chunk. [`BitPackedPatchedPlugin`] reads both and lifts +//! interior patches into a `Patched` array. mod bitpacked; mod patched; @@ -9,4 +14,7 @@ mod patched; #[cfg(test)] pub(crate) use bitpacked::BitPackedMetadata; pub use bitpacked::BitPackedPlugin; +#[cfg(test)] +pub(crate) use bitpacked::BitPackedV2Metadata; +pub use bitpacked::bitpacked_v2_id; pub(crate) use patched::BitPackedPatchedPlugin; diff --git a/encodings/fastlanes/src/bitpacking/plugin/patched.rs b/encodings/fastlanes/src/bitpacking/plugin/patched.rs index 5e73fb486f4..1593765313d 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/patched.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/patched.rs @@ -29,8 +29,11 @@ impl ArrayPlugin for BitPackedPatchedPlugin { fn id(&self) -> ArrayId { // We reuse the existing `BitPacked` ID so that we can take over its // deserialization pathway. - // TODO(joe): dedup method name - ArrayVTable::id(&BitPacked) + BitPackedPlugin.id() + } + + fn serialized_ids(&self) -> Vec { + BitPackedPlugin.serialized_ids() } fn serialize( @@ -38,7 +41,6 @@ impl ArrayPlugin for BitPackedPatchedPlugin { array: &ArrayRef, session: &VortexSession, ) -> VortexResult> { - // Both plugins share the same wire contract. BitPackedPlugin.serialize(array, session) } @@ -58,13 +60,13 @@ impl ArrayPlugin for BitPackedPatchedPlugin { let packed = bitpacked.packed().clone(); let ptype = bitpacked.dtype().as_ptype(); let validity = bitpacked.validity()?; - let bw = bitpacked.width_table().clone(); + let widths = bitpacked.width_table().clone(); let offsets = bitpacked.chunk_offsets().clone(); let len = bitpacked.len(); let offset = bitpacked.offset(); let bitpacked_without_patches = - BitPacked::try_new(packed, ptype, validity, None, bw, offsets, len, offset)? + BitPacked::try_new(packed, ptype, validity, None, widths, offsets, len, offset)? .into_array(); let patched = Patched::from_array_and_patches( diff --git a/encodings/fastlanes/src/bitpacking/serde_tests.rs b/encodings/fastlanes/src/bitpacking/serde_tests.rs deleted file mode 100644 index bd750697fc9..00000000000 --- a/encodings/fastlanes/src/bitpacking/serde_tests.rs +++ /dev/null @@ -1,83 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::LazyLock; - -use prost::Message; -use vortex_array::ArrayContext; -use vortex_array::ArrayId; -use vortex_array::ArrayRef; -use vortex_array::ArrayVTable; -use vortex_array::VortexSessionExecute; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::assert_arrays_eq; -use vortex_array::serde::SerializeOptions; -use vortex_array::serde::SerializedArray; -use vortex_array::session::ArraySessionExt; -use vortex_buffer::ByteBufferMut; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_session::VortexSession; -use vortex_session::registry::ReadContext; - -use crate::BitPacked; -use crate::BitPackedArray; -use crate::bitpacking::bitpack_compress::bitpack_to_best_bit_width; -use crate::bitpacking::plugin::BitPackedMetadata; - -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); - -fn serde_roundtrip(array: &BitPackedArray) -> VortexResult<(ArrayId, Vec, ArrayRef)> { - let array_ref = array.as_array(); - let serialization = SESSION - .array_serialize(array_ref)? - .ok_or_else(|| vortex_err!("BitPacked must serialize"))?; - let array_ctx = ArrayContext::empty(); - let buffers = array_ref.serialize(&array_ctx, &SESSION, &SerializeOptions::default())?; - let mut bytes = ByteBufferMut::empty(); - for buffer in buffers { - bytes.extend_from_slice(&buffer); - } - let read = SerializedArray::try_from(bytes.freeze())?.decode( - array_ref.dtype(), - array_ref.len(), - &ReadContext::new(array_ctx.to_ids()), - &SESSION, - )?; - Ok((serialization.serialized_id, serialization.metadata, read)) -} - -#[test] -fn uniform_widths_serialize_as_original_format() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values: Vec = (0..3000).map(|i| i % 128).collect(); - let packed = - bitpack_to_best_bit_width(&PrimitiveArray::from_iter(values.iter().copied()), &mut ctx)?; - assert_eq!(packed.as_array().children().len(), 2); - assert!( - SESSION - .array_serialize(packed.as_array())? - .ok_or_else(|| vortex_err!("must serialize"))? - .children - .is_empty() - ); - let (id, metadata, read) = serde_roundtrip(&packed)?; - assert_eq!(id, ArrayVTable::id(&BitPacked)); - let original = BitPackedMetadata { - bit_width: 7, - offset: 0, - patches: None, - } - .encode_to_vec(); - assert_eq!(metadata, original); - assert_arrays_eq!( - read, - PrimitiveArray::from_iter(values.iter().copied()), - &mut ctx - ); - Ok(()) -}