Merge pull request #1674 from quickwit-oss/u128_codec_header

add header with codec type for u128

fastfield_codecs/src/compact_space/mod.rs

@@ -456,6 +456,7 @@ impl CompactSpaceDecompressor {
 mod tests {
 
     use super::*;
+    use crate::serialize::U128Header;
     use crate::{open_u128, serialize_u128};
 
     #[test]
@@ -501,7 +502,8 @@ mod tests {
         assert_eq!(amplitude, 2);
     }
 
-    fn test_all(data: OwnedBytes, expected: &[u128]) {
+    fn test_all(mut data: OwnedBytes, expected: &[u128]) {
+        let _header = U128Header::deserialize(&mut data);
         let decompressor = CompactSpaceDecompressor::open(data).unwrap();
         for (idx, expected_val) in expected.iter().cloned().enumerate() {
             let val = decompressor.get(idx as u32);
@@ -556,7 +558,8 @@ mod tests {
             4_000_211_222u128,
             333u128,
         ];
-        let data = test_aux_vals(vals);
+        let mut data = test_aux_vals(vals);
+        let _header = U128Header::deserialize(&mut data);
         let decomp = CompactSpaceDecompressor::open(data).unwrap();
         let complete_range = 0..vals.len() as u32;
         for (pos, val) in vals.iter().enumerate() {
@@ -681,7 +684,8 @@ mod tests {
             4_000_211_222u128,
             333u128,
         ];
-        let data = test_aux_vals(vals);
+        let mut data = test_aux_vals(vals);
+        let _header = U128Header::deserialize(&mut data);
         let decomp = CompactSpaceDecompressor::open(data).unwrap();
         let complete_range = 0..vals.len() as u32;
         assert_eq!(

fastfield_codecs/src/lib.rs

@@ -25,7 +25,7 @@ use monotonic_mapping::{
     StrictlyMonotonicMappingToInternalBaseval, StrictlyMonotonicMappingToInternalGCDBaseval,
 };
 use ownedbytes::OwnedBytes;
-use serialize::Header;
+use serialize::{Header, U128Header};
 
 mod bitpacked;
 mod blockwise_linear;
@@ -92,10 +92,47 @@ impl FastFieldCodecType {
     }
 }
 
+#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
+#[repr(u8)]
+/// Available codecs to use to encode the u128 (via [`MonotonicallyMappableToU128`]) converted data.
+pub enum U128FastFieldCodecType {
+    /// This codec takes a large number space (u128) and reduces it to a compact number space, by
+    /// removing the holes.
+    CompactSpace = 1,
+}
+
+impl BinarySerializable for U128FastFieldCodecType {
+    fn serialize<W: Write>(&self, wrt: &mut W) -> io::Result<()> {
+        self.to_code().serialize(wrt)
+    }
+
+    fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
+        let code = u8::deserialize(reader)?;
+        let codec_type: Self = Self::from_code(code)
+            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Unknown code `{code}.`"))?;
+        Ok(codec_type)
+    }
+}
+
+impl U128FastFieldCodecType {
+    pub(crate) fn to_code(self) -> u8 {
+        self as u8
+    }
+
+    pub(crate) fn from_code(code: u8) -> Option<Self> {
+        match code {
+            1 => Some(Self::CompactSpace),
+            _ => None,
+        }
+    }
+}
+
 /// Returns the correct codec reader wrapped in the `Arc` for the data.
 pub fn open_u128<Item: MonotonicallyMappableToU128>(
-    bytes: OwnedBytes,
+    mut bytes: OwnedBytes,
 ) -> io::Result<Arc<dyn Column<Item>>> {
+    let header = U128Header::deserialize(&mut bytes)?;
+    assert_eq!(header.codec_type, U128FastFieldCodecType::CompactSpace);
     let reader = CompactSpaceDecompressor::open(bytes)?;
     let inverted: StrictlyMonotonicMappingInverter<StrictlyMonotonicMappingToInternal<Item>> =
         StrictlyMonotonicMappingToInternal::<Item>::new().into();

fastfield_codecs/src/serialize.rs

@@ -35,7 +35,7 @@ use crate::monotonic_mapping::{
 };
 use crate::{
     monotonic_map_column, Column, FastFieldCodec, FastFieldCodecType, MonotonicallyMappableToU64,
-    VecColumn, ALL_CODEC_TYPES,
+    U128FastFieldCodecType, VecColumn, ALL_CODEC_TYPES,
 };
 
 /// The normalized header gives some parameters after applying the following
@@ -98,6 +98,29 @@ impl Header {
     }
 }
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub(crate) struct U128Header {
+    pub num_vals: u32,
+    pub codec_type: U128FastFieldCodecType,
+}
+
+impl BinarySerializable for U128Header {
+    fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
+        VInt(self.num_vals as u64).serialize(writer)?;
+        self.codec_type.serialize(writer)?;
+        Ok(())
+    }
+
+    fn deserialize<R: io::Read>(reader: &mut R) -> io::Result<Self> {
+        let num_vals = VInt::deserialize(reader)?.0 as u32;
+        let codec_type = U128FastFieldCodecType::deserialize(reader)?;
+        Ok(U128Header {
+            num_vals,
+            codec_type,
+        })
+    }
+}
+
 pub fn normalize_column<C: Column>(
     from_column: C,
     min_value: u64,
@@ -167,7 +190,11 @@ pub fn serialize_u128<F: Fn() -> I, I: Iterator<Item = u128>>(
     num_vals: u32,
     output: &mut impl io::Write,
 ) -> io::Result<()> {
-    // TODO write header, to later support more codecs
+    let header = U128Header {
+        num_vals,
+        codec_type: U128FastFieldCodecType::CompactSpace,
+    };
+    header.serialize(output)?;
     let compressor = CompactSpaceCompressor::train_from(iter_gen(), num_vals);
     compressor.compress_into(iter_gen(), output).unwrap();
 
@@ -258,6 +285,18 @@ pub fn serialize_and_load<T: MonotonicallyMappableToU64 + Ord + Default>(
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_serialize_deserialize_u128_header() {
+        let original = U128Header {
+            num_vals: 11,
+            codec_type: U128FastFieldCodecType::CompactSpace,
+        };
+        let mut out = Vec::new();
+        original.serialize(&mut out).unwrap();
+        let restored = U128Header::deserialize(&mut &out[..]).unwrap();
+        assert_eq!(restored, original);
+    }
+
     #[test]
     fn test_serialize_deserialize() {
         let original = [1u64, 5u64, 10u64];