Fix TurboQuant metadata to be protobuf (#7301)

## Summary

Right now, the `TurboQuant` metadata is just a byte, this changes it to
a prost message.

## Testing

Small roundtrip tests.

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-tensor/src/encodings/turboquant/metadata.rs

@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Protobuf-backed metadata for TurboQuant encoding.
+
+use prost::Message;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_err;
+
+/// Serialized metadata for TurboQuant arrays.
+#[derive(Clone, PartialEq, Message)]
+pub(super) struct TurboQuantMetadata {
+    /// The number of bits per coordinate.
+    #[prost(uint32, required, tag = "1")]
+    bit_width: u32,
+}
+
+impl TurboQuantMetadata {
+    /// Creates metadata for the given bit width.
+    pub(super) fn new(bit_width: u8) -> Self {
+        Self {
+            bit_width: u32::from(bit_width),
+        }
+    }
+
+    /// Returns the validated TurboQuant bit width.
+    pub(super) fn bit_width(&self) -> VortexResult<u8> {
+        let bit_width = u8::try_from(self.bit_width).map_err(|_| {
+            vortex_err!(
+                "TurboQuant bit_width must fit into u8, got {}",
+                self.bit_width
+            )
+        })?;
+        vortex_ensure!(
+            bit_width <= 8,
+            "bit_width is expected to be between 0 and 8, got {bit_width}"
+        );
+
+        Ok(bit_width)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use prost::Message;
+    use rstest::rstest;
+    use vortex_error::VortexResult;
+
+    use super::TurboQuantMetadata;
+
+    #[rstest]
+    #[case(0)]
+    #[case(3)]
+    #[case(8)]
+    fn protobuf_metadata_roundtrip(#[case] bit_width: u8) -> VortexResult<()> {
+        let bytes = TurboQuantMetadata::new(bit_width).encode_to_vec();
+        assert_eq!(
+            TurboQuantMetadata::decode(bytes.as_slice())?.bit_width()?,
+            bit_width
+        );
+
+        Ok(())
+    }
+}

vortex-tensor/src/encodings/turboquant/mod.rs

@@ -97,6 +97,8 @@ pub use array::scheme::TurboQuantScheme;
 
 pub(crate) mod compute;
 
+mod metadata;
+
 mod vtable;
 pub use vtable::TurboQuant;
 pub use vtable::TurboQuantArray;

vortex-tensor/src/encodings/turboquant/vtable.rs

@@ -7,6 +7,7 @@ use std::hash::Hash;
 use std::hash::Hasher;
 use std::sync::Arc;
 
+use prost::Message;
 use vortex_array::Array;
 use vortex_array::ArrayEq;
 use vortex_array::ArrayHash;
@@ -40,6 +41,7 @@ use crate::encodings::turboquant::array::slots::Slot;
 use crate::encodings::turboquant::compute::rules::PARENT_KERNELS;
 use crate::encodings::turboquant::compute::rules::RULES;
 use crate::encodings::turboquant::decompress::execute_decompress;
+use crate::encodings::turboquant::metadata::TurboQuantMetadata;
 use crate::utils::tensor_element_ptype;
 use crate::utils::tensor_list_size;
 use crate::vector::Vector;
@@ -197,7 +199,9 @@ impl VTable for TurboQuant {
     }
 
     fn serialize(array: ArrayView<'_, Self>) -> VortexResult<Option<Vec<u8>>> {
-        Ok(Some(vec![array.bit_width]))
+        Ok(Some(
+            TurboQuantMetadata::new(array.bit_width).encode_to_vec(),
+        ))
     }
 
     fn deserialize(
@@ -209,19 +213,8 @@ impl VTable for TurboQuant {
         children: &dyn ArrayChildren,
         _session: &VortexSession,
     ) -> VortexResult<ArrayParts<Self>> {
-        vortex_ensure_eq!(
-            metadata.len(),
-            1,
-            "TurboQuant metadata must be exactly 1 byte, got {}",
-            metadata.len()
-        );
-        vortex_ensure!(
-            metadata[0] <= 8,
-            "bit_width is expected to be between 0 and 8, got {}",
-            metadata[0]
-        );
-
-        let bit_width = metadata[0];
+        let metadata = TurboQuantMetadata::decode(metadata)?;
+        let bit_width = metadata.bit_width()?;
 
         // bit_width == 0 is only valid for degenerate (empty) arrays. A non-empty array with
         // bit_width == 0 would have zero centroids while codes reference centroid indices.