refactor: isolate bitpacked v1 serialization

Signed-off-by: "Matt Katz" <mhkatz97@gmail.com>

Author: mhk197

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;

encodings/fastlanes/src/bitpacking/plugin.rs

@@ -1,11 +1,9 @@
 // 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.
+//! Serialization plugins for the frozen bit-packed format and its external-patches adapter.
 
+use prost::Message;
 use vortex_array::Array;
 use vortex_array::ArrayDeserialization;
 use vortex_array::ArrayId;
@@ -18,32 +16,41 @@ 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 crate::BitPacked;
 use crate::BitPackedArrayExt;
+use crate::bitpacking::vtable::BitPackedMetadata;
 
-/// Custom deserialization plugin that converts a BitPacked array with interior
-/// Patches into a PatchedArray holding a BitPacked array.
+/// Serialization boundary for the frozen `fastlanes.bitpacked` wire format.
 #[derive(Debug, Clone)]
-pub(crate) struct BitPackedPatchedPlugin;
+pub struct BitPackedPlugin;
 
-impl ArrayPlugin for BitPackedPatchedPlugin {
+impl ArrayPlugin for BitPackedPlugin {
     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)
     }
 
     fn serialize(
         &self,
         array: &ArrayRef,
-        session: &VortexSession,
+        _session: &VortexSession,
     ) -> VortexResult<Option<ArraySerialization>> {
-        // delegate to BitPacked VTable for serialization
-        ArrayPlugin::serialize(&BitPacked, array, session)
+        let view = array.as_::<BitPacked>();
+        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(
@@ -54,18 +61,50 @@ impl ArrayPlugin for BitPackedPatchedPlugin {
         vortex_ensure!(
             parts.serialized_id == self.id(),
             "BitPacked plugin does not recognize serialized ID {}",
-            parts.serialized_id,
+            parts.serialized_id
         );
-        let bitpacked = Array::<BitPacked>::try_from_parts(ArrayVTable::deserialize(
+        Ok(Array::<BitPacked>::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"))?;
+        )?)?
+        .into_array())
+    }
+}
+
+/// Custom deserialization plugin that converts a BitPacked array with interior
+/// Patches into a PatchedArray holding a BitPacked array.
+#[derive(Debug, Clone)]
+pub(crate) struct BitPackedPatchedPlugin;
+
+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)
+    }
+
+    fn serialize(
+        &self,
+        array: &ArrayRef,
+        session: &VortexSession,
+    ) -> VortexResult<Option<ArraySerialization>> {
+        // Both plugins share the same wire contract.
+        BitPackedPlugin.serialize(array, session)
+    }
+
+    fn deserialize(
+        &self,
+        parts: ArrayDeserialization<'_>,
+        session: &VortexSession,
+    ) -> VortexResult<ArrayRef> {
+        let bitpacked = BitPackedPlugin.deserialize(parts, session)?;
+        let bitpacked = bitpacked.as_::<BitPacked>().into_owned();
 
         // Create a new BitPackedArray without the interior patches installed.
         let Some(patches) = bitpacked.patches() else {

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::vtable::BitPackedMetadata;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = vortex_array::array_session();
+    crate::initialize(&session);
+    session
+});
+
+fn serde_roundtrip(array: &BitPackedArray) -> VortexResult<(ArrayId, Vec<u8>, 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<u32> = (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(())
+}

encodings/fastlanes/src/bitpacking/vtable/mod.rs

@@ -161,20 +161,10 @@ impl VTable for BitPacked {
     }
 
     fn serialize(
-        array: ArrayView<'_, Self>,
+        _array: ArrayView<'_, Self>,
         _session: &VortexSession,
     ) -> VortexResult<Option<Vec<u8>>> {
-        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(
@@ -187,71 +177,24 @@ impl VTable for BitPacked {
         _session: &VortexSession,
     ) -> VortexResult<ArrayParts<Self>> {
         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 packed = single_buffer(buffers)?;
+        let offset = offset_from_metadata(metadata.offset)?;
+        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 (patches, validity, _) =
+            deserialize_children(children, metadata.patches, dtype, len, 0)?;
 
         let slots = {
-            let mut s = ArraySlots::with_capacity(4);
+            let mut s = ArraySlots::with_capacity(BitPackedSlots::COUNT);
             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
-                )
-            })?,
-        )?;
+        let data = BitPackedData::try_new(packed, patches, bit_width, offset)?;
         Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
     }
 
@@ -299,6 +242,71 @@ impl VTable for BitPacked {
     }
 }
 
+/// The single packed buffer of a serialized bit-packed array.
+pub(crate) fn single_buffer(buffers: &[BufferHandle]) -> VortexResult<BufferHandle> {
+    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`.
+pub(crate) fn offset_from_metadata(offset: u32) -> VortexResult<u16> {
+    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.
+pub(crate) fn deserialize_children(
+    children: &dyn ArrayChildren,
+    patches: Option<PatchesMetadata>,
+    dtype: &DType,
+    len: usize,
+    trailing: usize,
+) -> VortexResult<(Option<Patches>, 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),
+    ))
+}
+
 #[derive(Clone, Debug)]
 pub struct BitPacked;