diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index efa0677a91e..cf556df2780 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -17,9 +17,13 @@ mod plugin; mod vtable; pub(crate) use plugin::BitPackedPatchedPlugin; +pub use plugin::BitPackedPlugin; pub use vtable::BitPacked; pub use vtable::BitPackedArray; pub(crate) fn initialize(session: &vortex_session::VortexSession) { vtable::initialize(session); } + +#[cfg(test)] +mod serde_tests; diff --git a/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs b/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs new file mode 100644 index 00000000000..4cfc08011fa --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serialization plugin for the frozen bit-packed wire format. + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::ArraySlots; +use vortex_array::ArrayVTable; +use vortex_array::IntoArray; +use vortex_array::patches::Patches; +use vortex_array::patches::PatchesData; +use vortex_array::patches::PatchesMetadata; +use vortex_array::validity::Validity; +use vortex_array::vtable::validity_to_child; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::BitPacked; +use crate::BitPackedArrayExt; +use crate::BitPackedData; + +/// Metadata of the frozen `fastlanes.bitpacked` wire format. +#[derive(Clone, prost::Message)] +pub(crate) struct BitPackedMetadata { + #[prost(uint32, tag = "1")] + pub(crate) bit_width: u32, + #[prost(uint32, tag = "2")] + pub(crate) offset: u32, // must be <1024 + #[prost(message, optional, tag = "3")] + pub(crate) patches: Option, +} + +/// Serialization boundary for the frozen `fastlanes.bitpacked` wire format. +#[derive(Debug, Clone)] +pub struct BitPackedPlugin; + +impl ArrayPlugin for BitPackedPlugin { + fn id(&self) -> ArrayId { + ArrayVTable::id(&BitPacked) + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + self.id() == array.encoding_id(), + "array plugin {} cannot serialize in-memory array {}", + self.id(), + array.encoding_id(), + ); + let view = array.as_::(); + let metadata = BitPackedMetadata { + bit_width: view.bit_width() as u32, + offset: view.offset() as u32, + patches: view + .patches() + .map(|p| p.to_metadata(view.len(), view.dtype())) + .transpose()?, + } + .encode_to_vec(); + Ok(Some(ArraySerialization::from_array( + self.id(), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure!( + self.id() == parts.serialized_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 slots = { + let mut s = ArraySlots::with_capacity(4); + PatchesData::push_slots(&mut s, patches.as_ref()); + s.push(validity_to_child(&validity, len)); + s + }; + let data = BitPackedData::try_new( + packed, + patches, + u8::try_from(metadata.bit_width).map_err(|_| { + vortex_err!( + "BitPackedMetadata bit_width {} does not fit in u8", + metadata.bit_width + ) + })?, + u16::try_from(metadata.offset).map_err(|_| { + vortex_err!( + "BitPackedMetadata offset {} does not fit in u16", + metadata.offset + ) + })?, + )?; + Ok(Array::::try_from_parts( + ArrayParts::new(BitPacked, dtype.clone(), len, data).with_slots(slots), + )? + .into_array()) + } +} + +#[cfg(test)] +mod tests { + use prost::Message; + use rstest::rstest; + use vortex_array::ArrayDeserialization; + use vortex_array::ArrayPlugin; + use vortex_array::ArrayVTable; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::buffer::BufferHandle; + use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::BitPackedMetadata; + use super::BitPackedPlugin; + use crate::BitPacked; + use crate::BitPackedData; + + #[test] + fn serialize_rejects_other_encodings() { + let session = vortex_array::array_session(); + let values = PrimitiveArray::from_iter([1u32, 2, 3]).into_array(); + assert!(BitPackedPlugin.serialize(&values, &session).is_err()); + } + + #[rstest] + #[case(2, u32::MAX, u32::MAX, "Expected 0 or 1 children")] + #[case(0, u32::MAX, u32::MAX, "bit_width")] + #[case(0, 0, u32::MAX, "offset")] + fn invalid_inputs_preserve_validation_order( + #[case] num_children: usize, + #[case] bit_width: u32, + #[case] offset: u32, + #[case] expected: &str, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let child = PrimitiveArray::from_iter([0u32]).into_array(); + let metadata = BitPackedMetadata { + bit_width, + offset, + patches: None, + } + .encode_to_vec(); + let buffers = [BufferHandle::new_host(ByteBuffer::empty())]; + let children = vec![child.clone(); num_children]; + let error = BitPackedPlugin + .deserialize( + ArrayDeserialization::new( + BitPackedPlugin.id(), + child.dtype(), + 0, + &metadata, + &buffers, + &children, + ), + &session, + ) + .err() + .ok_or_else(|| vortex_err!("Invalid inputs must be rejected"))?; + assert!(error.to_string().contains(expected), "{error}"); + Ok(()) + } + + #[test] + fn serde_requires_plugin() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let values = + PrimitiveArray::from_option_iter([Some(1u32), None, Some(511), Some(7)]).into_array(); + let packed = BitPackedData::encode(&values, 3, &mut ctx)?; + let array = packed.as_array(); + let serialized = BitPackedPlugin + .serialize(array, &session)? + .ok_or_else(|| vortex_err!("BitPacked must serialize"))?; + let buffers = serialized + .buffers + .iter() + .cloned() + .map(BufferHandle::new_host) + .collect::>(); + + assert!(ArrayVTable::serialize(packed.as_view(), &session).is_err()); + assert!( + ArrayVTable::deserialize( + &BitPacked, + array.dtype(), + array.len(), + &serialized.metadata, + &buffers, + &serialized.children, + &session, + ) + .is_err() + ); + + let decoded = BitPackedPlugin.deserialize( + ArrayDeserialization::new( + BitPackedPlugin.id(), + array.dtype(), + array.len(), + &serialized.metadata, + &buffers, + &serialized.children, + ), + &session, + )?; + assert_arrays_eq!(decoded, values, &mut ctx); + Ok(()) + } +} diff --git a/encodings/fastlanes/src/bitpacking/plugin/mod.rs b/encodings/fastlanes/src/bitpacking/plugin/mod.rs new file mode 100644 index 00000000000..d97594f7c26 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/plugin/mod.rs @@ -0,0 +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. + +mod bitpacked; +mod patched; + +#[cfg(test)] +pub(crate) use bitpacked::BitPackedMetadata; +pub use bitpacked::BitPackedPlugin; +pub(crate) use patched::BitPackedPatchedPlugin; diff --git a/encodings/fastlanes/src/bitpacking/plugin.rs b/encodings/fastlanes/src/bitpacking/plugin/patched.rs similarity index 87% rename from encodings/fastlanes/src/bitpacking/plugin.rs rename to encodings/fastlanes/src/bitpacking/plugin/patched.rs index 3ff07db7e5c..9a42a60a7a5 100644 --- a/encodings/fastlanes/src/bitpacking/plugin.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/patched.rs @@ -1,12 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A custom [`ArrayPlugin`] that lets you load in and deserialize a `BitPacked` array with interior -//! patches as a `PatchedArray` that wraps a patchless `BitPacked` array. -//! -//! This enables zero-cost backward compatibility with previously written datasets. +//! Deserialization adapter that lifts internal bit-packed patches into a `Patched` array. -use vortex_array::Array; use vortex_array::ArrayDeserialization; use vortex_array::ArrayId; use vortex_array::ArrayPlugin; @@ -17,10 +13,9 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Patched; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_session::VortexSession; +use super::BitPackedPlugin; use crate::BitPacked; use crate::BitPackedArrayExt; @@ -42,8 +37,8 @@ impl ArrayPlugin for BitPackedPatchedPlugin { array: &ArrayRef, session: &VortexSession, ) -> VortexResult> { - // delegate to BitPacked VTable for serialization - ArrayPlugin::serialize(&BitPacked, array, session) + // Both plugins share the same wire contract. + BitPackedPlugin.serialize(array, session) } fn deserialize( @@ -51,21 +46,8 @@ impl ArrayPlugin for BitPackedPatchedPlugin { parts: ArrayDeserialization<'_>, session: &VortexSession, ) -> VortexResult { - vortex_ensure!( - parts.serialized_id == self.id(), - "BitPacked plugin does not recognize serialized ID {}", - parts.serialized_id, - ); - let bitpacked = Array::::try_from_parts(ArrayVTable::deserialize( - &BitPacked, - parts.dtype, - parts.len, - parts.metadata, - parts.buffers, - parts.children, - session, - )?) - .map_err(|_| vortex_err!("BitPacked plugin should only deserialize fastlanes.bitpacked"))?; + let bitpacked = BitPackedPlugin.deserialize(parts, session)?; + let bitpacked = bitpacked.as_::().into_owned(); // Create a new BitPackedArray without the interior patches installed. let Some(patches) = bitpacked.patches() else { diff --git a/encodings/fastlanes/src/bitpacking/serde_tests.rs b/encodings/fastlanes/src/bitpacking/serde_tests.rs new file mode 100644 index 00000000000..7c8a8d3d952 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/serde_tests.rs @@ -0,0 +1,83 @@ +// 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!(packed.as_array().children().is_empty()); + 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(()) +} diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 68fbf1b41d3..aa18bd289c6 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -4,7 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use prost::Message; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; @@ -24,7 +23,6 @@ use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; use vortex_array::patches::Patches; use vortex_array::patches::PatchesData; -use vortex_array::patches::PatchesMetadata; use vortex_array::require_patches; use vortex_array::require_validity; use vortex_array::serde::ArrayChildren; @@ -36,7 +34,6 @@ use vortex_error::VortexExpect; 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_session::VortexSession; use vortex_session::registry::CachedId; @@ -62,16 +59,6 @@ pub(crate) fn initialize(session: &VortexSession) { kernels::initialize(session); } -#[derive(Clone, prost::Message)] -pub struct BitPackedMetadata { - #[prost(uint32, tag = "1")] - pub(crate) bit_width: u32, - #[prost(uint32, tag = "2")] - pub(crate) offset: u32, // must be <1024 - #[prost(message, optional, tag = "3")] - pub(crate) patches: Option, -} - impl ArrayHash for BitPackedData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.offset.hash(state); @@ -161,98 +148,22 @@ impl VTable for BitPacked { } fn serialize( - array: ArrayView<'_, Self>, + _array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - Ok(Some( - BitPackedMetadata { - bit_width: array.bit_width() as u32, - offset: array.offset() as u32, - patches: array - .patches() - .map(|p| p.to_metadata(array.len(), array.dtype())) - .transpose()?, - } - .encode_to_vec(), - )) + vortex_bail!("BitPacked serialization requires BitPackedPlugin") } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], - buffers: &[BufferHandle], - children: &dyn ArrayChildren, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - 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 slots = { - let mut s = ArraySlots::with_capacity(4); - PatchesData::push_slots(&mut s, patches.as_ref()); - s.push(validity_to_child(&validity, len)); - s - }; - let data = BitPackedData::try_new( - packed, - patches, - u8::try_from(metadata.bit_width).map_err(|_| { - vortex_err!( - "BitPackedMetadata bit_width {} does not fit in u8", - metadata.bit_width - ) - })?, - u16::try_from(metadata.offset).map_err(|_| { - vortex_err!( - "BitPackedMetadata offset {} does not fit in u16", - metadata.offset - ) - })?, - )?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + vortex_bail!("BitPacked deserialization requires BitPackedPlugin") } fn append_to_builder( diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 43d83c6fc7f..ee7a6838883 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -84,7 +84,7 @@ pub fn initialize(session: &VortexSession) { if use_experimental_patches() { session.arrays().register(BitPackedPatchedPlugin); } else { - session.arrays().register(BitPacked); + session.arrays().register(BitPackedPlugin); } session.arrays().register(Delta); session.arrays().register(FoR);