Move bool array over to ProstMetadata (#3052)

Part of #3051

Author: gatesn

Cargo.lock

@@ -42,6 +42,8 @@ num_enum = { workspace = true }
 parking_lot = { workspace = true }
 paste = { workspace = true }
 pin-project = { workspace = true }
+prost = { workspace = true }
+prost-types = { workspace = true }
 rand = { workspace = true }
 rkyv = { workspace = true }
 rstest = { workspace = true, optional = true }
@@ -54,6 +56,7 @@ vortex-dtype = { workspace = true, features = ["arrow", "rkyv", "serde"] }
 vortex-error = { workspace = true, features = [
     "flatbuffers",
     "flexbuffers",
+    "prost",
     "rancor",
 ] }
 vortex-flatbuffers = { workspace = true, features = ["array"] }

vortex-array/Cargo.toml

@@ -12,7 +12,7 @@ use crate::stats::{ArrayStats, StatsSetRef};
 use crate::validity::Validity;
 use crate::variants::BoolArrayTrait;
 use crate::vtable::VTableRef;
-use crate::{ArrayImpl, ArrayRef, ArrayStatisticsImpl, Canonical, Encoding, RkyvMetadata};
+use crate::{ArrayImpl, ArrayRef, ArrayStatisticsImpl, Canonical, Encoding, ProstMetadata};
 
 #[derive(Clone, Debug)]
 pub struct BoolArray {
@@ -26,7 +26,7 @@ pub struct BoolArray {
 pub struct BoolEncoding;
 impl Encoding for BoolEncoding {
     type Array = BoolArray;
-    type Metadata = RkyvMetadata<BoolMetadata>;
+    type Metadata = ProstMetadata<BoolMetadata>;
 }
 
 impl BoolArray {

vortex-array/src/arrays/bool/array.rs

@@ -9,9 +9,16 @@ use crate::validity::Validity;
 use crate::vtable::EncodingVTable;
 use crate::{
     Array, ArrayBufferVisitor, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl,
-    DeserializeMetadata, EncodingId, RkyvMetadata,
+    DeserializeMetadata, Encoding, EncodingId, ProstMetadata,
 };
 
+#[derive(prost::Message)]
+pub struct BoolMetadata {
+    // The offset in bits must be <8
+    #[prost(uint32, tag = "1")]
+    pub offset: u32,
+}
+
 impl EncodingVTable for BoolEncoding {
     fn id(&self) -> EncodingId {
         EncodingId::new_ref("vortex.bool")
@@ -24,7 +31,7 @@ impl EncodingVTable for BoolEncoding {
         dtype: DType,
         len: usize,
     ) -> VortexResult<ArrayRef> {
-        let metadata = RkyvMetadata::<BoolMetadata>::deserialize(parts.metadata())?;
+        let metadata = <Self as Encoding>::Metadata::deserialize(parts.metadata())?;
 
         if parts.nbuffers() != 1 {
             vortex_bail!("Expected 1 buffer, got {}", parts.nbuffers());
@@ -48,13 +55,7 @@ impl EncodingVTable for BoolEncoding {
     }
 }
 
-#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
-pub struct BoolMetadata {
-    // We know the offset in bits must be <8
-    offset: u8,
-}
-
-impl ArrayVisitorImpl<RkyvMetadata<BoolMetadata>> for BoolArray {
+impl ArrayVisitorImpl<ProstMetadata<BoolMetadata>> for BoolArray {
     fn _visit_buffers(&self, visitor: &mut dyn ArrayBufferVisitor) {
         visitor.visit_buffer(&ByteBuffer::from_arrow_buffer(
             self.boolean_buffer().clone().into_inner(),
@@ -66,11 +67,11 @@ impl ArrayVisitorImpl<RkyvMetadata<BoolMetadata>> for BoolArray {
         visitor.visit_validity(&self.validity, self.len());
     }
 
-    fn _metadata(&self) -> RkyvMetadata<BoolMetadata> {
+    fn _metadata(&self) -> ProstMetadata<BoolMetadata> {
         let bit_offset = self.boolean_buffer().offset();
         assert!(bit_offset < 8, "Offset must be <8, got {}", bit_offset);
-        RkyvMetadata(BoolMetadata {
-            offset: u8::try_from(bit_offset).vortex_expect("checked"),
+        ProstMetadata(BoolMetadata {
+            offset: u32::try_from(bit_offset).vortex_expect("checked"),
         })
     }
 }

vortex-array/src/arrays/bool/serde.rs

@@ -132,6 +132,46 @@ where
     }
 }
 
+/// A utility wrapper for Prost metadata serialization.
+pub struct ProstMetadata<M>(pub M);
+
+impl<M: Debug> Debug for ProstMetadata<M> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<M> SerializeMetadata for ProstMetadata<M>
+where
+    M: prost::Message,
+{
+    fn serialize(&self) -> Option<Vec<u8>> {
+        Some(self.0.encode_to_vec())
+    }
+}
+
+impl<M> DeserializeMetadata for ProstMetadata<M>
+where
+    M: Debug,
+    M: prost::Message + Default,
+{
+    type Output = M;
+
+    fn deserialize(metadata: Option<&[u8]>) -> VortexResult<Self::Output> {
+        let bytes =
+            metadata.ok_or_else(|| vortex_err!("Prost metadata requires metadata bytes"))?;
+        Ok(M::decode(bytes)?)
+    }
+
+    #[allow(clippy::use_debug)]
+    fn format(metadata: Option<&[u8]>, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match Self::deserialize(metadata) {
+            Ok(m) => write!(f, "{:?}", m),
+            Err(_) => write!(f, "Failed to deserialize metadata"),
+        }
+    }
+}
+
 /// A utility wrapper for automating the serialization of metadata using [serde](docs.rs/serde) into [flexbuffers](https://docs.rs/flexbuffers/latest/flexbuffers/).
 pub struct SerdeMetadata<M>(pub M);