Make doc store binary format more compact (#903)

CHANGELOG.md

@@ -52,6 +52,7 @@ Tantivy 0.26 (Unreleased)
 - Skip column traversal in `RangeDocSet` when query range does not overlap with column bounds [#2783](https://github.com/quickwit-oss/tantivy/pull/2783)(@ChangRui-Ryan)
 - Speed up exclude queries by supporting multiple excluded `DocSet`s without intermediate union [#2825](https://github.com/quickwit-oss/tantivy/pull/2825)(@PSeitz)
 - Improve union performance for non-score unions with `fill_buffer` and optimized `TinySet` [#2863](https://github.com/quickwit-oss/tantivy/pull/2863)(@PSeitz)
+- Make the doc store binary format more compact by storing field ids and integers as variable-length integers (`VInt`, with zig-zag encoding for signed integers). This introduces doc store format `V3` while remaining backwards compatible with older segments [#903](https://github.com/quickwit-oss/tantivy/issues/903)
 
 Tantivy 0.25
 ================================

common/src/lib.rs

@@ -23,6 +23,7 @@ pub use ownedbytes::{OwnedBytes, StableDeref};
 pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize};
 pub use vint::{
     VInt, VIntU128, read_u32_vint, read_u32_vint_no_advance, serialize_vint_u32, write_u32_vint,
+    zig_zag_decode, zig_zag_encode,
 };
 pub use writer::{AntiCallToken, CountingWriter, TerminatingWrite};
 

common/src/vint.rs

@@ -159,6 +159,22 @@ pub fn write_u32_vint<W: io::Write + ?Sized>(val: u32, writer: &mut W) -> io::Re
     writer.write_all(data)
 }
 
+/// Zig-zag encodes a signed integer into an unsigned integer.
+///
+/// Small-magnitude values (whether positive or negative) map to small
+/// unsigned values, which makes them cheap to encode as a [`VInt`].
+/// This is the same encoding used by protobuf and by the columnar crate.
+#[inline]
+pub fn zig_zag_encode(n: i64) -> u64 {
+    ((n << 1) ^ (n >> 63)) as u64
+}
+
+/// Reverses [`zig_zag_encode`].
+#[inline]
+pub fn zig_zag_decode(n: u64) -> i64 {
+    ((n >> 1) as i64) ^ (-((n & 1) as i64))
+}
+
 impl VInt {
     pub fn val(&self) -> u64 {
         self.0
@@ -226,7 +242,32 @@ impl BinarySerializable for VInt {
 #[cfg(test)]
 mod tests {
 
-    use super::{BinarySerializable, VInt, serialize_vint_u32};
+    use super::{BinarySerializable, VInt, serialize_vint_u32, zig_zag_decode, zig_zag_encode};
+
+    #[test]
+    fn test_zig_zag_roundtrip() {
+        for val in [
+            0,
+            1,
+            -1,
+            2,
+            -2,
+            42,
+            -42,
+            i64::MAX,
+            i64::MIN,
+            i64::MAX - 1,
+            i64::MIN + 1,
+        ] {
+            assert_eq!(zig_zag_decode(zig_zag_encode(val)), val);
+        }
+        // Small-magnitude values (positive or negative) encode to small unsigned
+        // values, which is the whole point of the transform.
+        assert_eq!(zig_zag_encode(0), 0);
+        assert_eq!(zig_zag_encode(-1), 1);
+        assert_eq!(zig_zag_encode(1), 2);
+        assert_eq!(zig_zag_encode(-2), 3);
+    }
 
     fn aux_test_vint(val: u64) {
         let mut v = [14u8; 10];

src/schema/document/de.rs

@@ -16,13 +16,13 @@ use std::net::Ipv6Addr;
 use std::sync::Arc;
 
 use columnar::MonotonicallyMappableToU128;
-use common::{u64_to_f64, BinarySerializable, DateTime, VInt};
+use common::{u64_to_f64, zig_zag_decode, BinarySerializable, DateTime, VInt};
 
 use super::se::BinaryObjectSerializer;
 use super::{OwnedValue, Value};
 use crate::schema::document::type_codes;
 use crate::schema::{Facet, Field};
-use crate::store::DocStoreVersion;
+use crate::store::{DocStoreVersion, DOC_STORE_VERSION};
 use crate::tokenizer::PreTokenizedString;
 
 #[derive(Debug, thiserror::Error, Clone)]
@@ -337,7 +337,14 @@ where R: Read
             return Ok(None);
         }
 
-        let field = Field::deserialize(self.reader).map_err(DeserializeError::from)?;
+        // Starting with `V3`, field ids are stored as a variable-length integer
+        // rather than a fixed 4-byte `u32`.
+        let field = if self.doc_store_version >= DocStoreVersion::V3 {
+            let field_id = VInt::deserialize(self.reader)?.val() as u32;
+            Field::from_field_id(field_id)
+        } else {
+            Field::deserialize(self.reader).map_err(DeserializeError::from)?
+        };
         let deserializer =
             BinaryValueDeserializer::from_reader(self.reader, self.doc_store_version)?;
         let value = V::deserialize(deserializer)?;
@@ -438,12 +445,26 @@ where R: Read
 
     fn deserialize_u64(self) -> Result<u64, DeserializeError> {
         self.validate_type(ValueType::U64)?;
-        <u64 as BinarySerializable>::deserialize(self.reader).map_err(DeserializeError::from)
+        if self.doc_store_version >= DocStoreVersion::V3 {
+            // Stored as a variable-length integer.
+            VInt::deserialize(self.reader)
+                .map(|vint| vint.val())
+                .map_err(DeserializeError::from)
+        } else {
+            <u64 as BinarySerializable>::deserialize(self.reader).map_err(DeserializeError::from)
+        }
     }
 
     fn deserialize_i64(self) -> Result<i64, DeserializeError> {
         self.validate_type(ValueType::I64)?;
-        <i64 as BinarySerializable>::deserialize(self.reader).map_err(DeserializeError::from)
+        if self.doc_store_version >= DocStoreVersion::V3 {
+            // Stored as a zig-zag encoded variable-length integer.
+            VInt::deserialize(self.reader)
+                .map(|vint| zig_zag_decode(vint.val()))
+                .map_err(DeserializeError::from)
+        } else {
+            <i64 as BinarySerializable>::deserialize(self.reader).map_err(DeserializeError::from)
+        }
     }
 
     fn deserialize_f64(self) -> Result<f64, DeserializeError> {
@@ -460,7 +481,7 @@ where R: Read
                 let timestamp_micros = <i64 as BinarySerializable>::deserialize(self.reader)?;
                 Ok(DateTime::from_timestamp_micros(timestamp_micros))
             }
-            DocStoreVersion::V2 => {
+            DocStoreVersion::V2 | DocStoreVersion::V3 => {
                 let timestamp_nanos = <i64 as BinarySerializable>::deserialize(self.reader)?;
                 Ok(DateTime::from_timestamp_nanos(timestamp_nanos))
             }
@@ -565,8 +586,9 @@ where R: Read
 
                 let out_rc = std::rc::Rc::new(out);
                 let mut slice: &[u8] = &out_rc;
-                let access =
-                    BinaryObjectDeserializer::from_reader(&mut slice, self.doc_store_version)?;
+                // `out` was just written by the current serializer, so it uses
+                // the latest doc store format regardless of the segment version.
+                let access = BinaryObjectDeserializer::from_reader(&mut slice, DOC_STORE_VERSION)?;
 
                 visitor.visit_object(access)
             }

src/schema/document/se.rs

@@ -3,7 +3,7 @@ use std::io;
 use std::io::Write;
 
 use columnar::MonotonicallyMappableToU128;
-use common::{f64_to_u64, BinarySerializable, VInt};
+use common::{f64_to_u64, zig_zag_encode, BinarySerializable, VInt};
 
 use super::{OwnedValue, ReferenceValueLeaf};
 use crate::schema::document::{type_codes, Document, ReferenceValue, Value};
@@ -37,7 +37,9 @@ where W: Write
 
         VInt(num_field_values as u64).serialize(self.writer)?;
         for (field, value_access) in stored_field_values() {
-            field.serialize(self.writer)?;
+            // Field ids are small and are stored as a variable-length integer
+            // (doc store `V3`+) rather than a fixed 4-byte `u32`.
+            VInt(field.field_id() as u64).serialize(self.writer)?;
 
             let mut serializer = BinaryValueSerializer::new(self.writer);
             match value_access.as_value() {
@@ -103,10 +105,13 @@ where W: Write
                     self.serialize_with_type_code(type_codes::TEXT_CODE, &Cow::Borrowed(val))
                 }
                 ReferenceValueLeaf::U64(val) => {
-                    self.serialize_with_type_code(type_codes::U64_CODE, &val)
+                    // Stored as a variable-length integer (doc store `V3`+).
+                    self.serialize_with_type_code(type_codes::U64_CODE, &VInt(val))
                 }
                 ReferenceValueLeaf::I64(val) => {
-                    self.serialize_with_type_code(type_codes::I64_CODE, &val)
+                    // Zig-zag encoded then stored as a variable-length integer
+                    // (doc store `V3`+) so small-magnitude negatives stay small.
+                    self.serialize_with_type_code(type_codes::I64_CODE, &VInt(zig_zag_encode(val)))
                 }
                 ReferenceValueLeaf::F64(val) => {
                     self.serialize_with_type_code(type_codes::F64_CODE, &f64_to_u64(val))
@@ -371,7 +376,7 @@ mod tests {
 
         let result = serialize_value(ReferenceValueLeaf::U64(123).into());
         let expected = binary_repr!(
-            type_codes::U64_CODE => 123u64,
+            type_codes::U64_CODE => VInt(123u64),
         );
         assert_eq!(
             result, expected,
@@ -380,7 +385,7 @@ mod tests {
 
         let result = serialize_value(ReferenceValueLeaf::I64(-123).into());
         let expected = binary_repr!(
-            type_codes::I64_CODE => -123i64,
+            type_codes::I64_CODE => VInt(zig_zag_encode(-123i64)),
         );
         assert_eq!(
             result, expected,
@@ -483,7 +488,7 @@ mod tests {
             length elements.len(),
             type_codes::NULL_CODE => (),
             type_codes::TEXT_CODE => String::from("Hello, world"),
-            type_codes::I64_CODE => 12345i64,
+            type_codes::I64_CODE => VInt(zig_zag_encode(12345i64)),
         );
         assert_eq!(
             result, expected,
@@ -601,7 +606,7 @@ mod tests {
             collection type_codes::OBJECT_CODE,
             length 4,   // 2 keys, 2 values
             type_codes::TEXT_CODE => String::from("inner-1"),
-            type_codes::I64_CODE => -123i64,
+            type_codes::I64_CODE => VInt(zig_zag_encode(-123i64)),
             type_codes::TEXT_CODE => String::from("inner-2"),
             type_codes::TEXT_CODE => String::from("bobby of the sea 2"),
         );
@@ -700,11 +705,15 @@ mod tests {
 
         let result = serialize_doc(&document, &schema);
         let mut expected = expected_doc_data!(length document.len());
-        name.serialize(&mut expected).unwrap();
+        VInt(name.field_id() as u64)
+            .serialize(&mut expected)
+            .unwrap();
         expected
             .extend_from_slice(&binary_repr!(type_codes::TEXT_CODE => String::from("ChillFish8")));
-        age.serialize(&mut expected).unwrap();
-        expected.extend_from_slice(&binary_repr!(type_codes::U64_CODE => 20u64));
+        VInt(age.field_id() as u64)
+            .serialize(&mut expected)
+            .unwrap();
+        expected.extend_from_slice(&binary_repr!(type_codes::U64_CODE => VInt(20u64)));
         assert_eq!(
             result, expected,
             "Expected serialized document to match the binary representation"
@@ -722,7 +731,9 @@ mod tests {
 
         let result = serialize_doc(&document, &schema);
         let mut expected = expected_doc_data!(length 1);
-        name.serialize(&mut expected).unwrap();
+        VInt(name.field_id() as u64)
+            .serialize(&mut expected)
+            .unwrap();
         expected
             .extend_from_slice(&binary_repr!(type_codes::TEXT_CODE => String::from("ChillFish8")));
         assert_eq!(

src/store/mod.rs

@@ -44,7 +44,7 @@ pub use self::writer::StoreWriter;
 mod store_compressor;
 
 /// Doc store version in footer to handle format changes.
-pub(crate) const DOC_STORE_VERSION: DocStoreVersion = DocStoreVersion::V2;
+pub(crate) const DOC_STORE_VERSION: DocStoreVersion = DocStoreVersion::V3;
 
 #[cfg(feature = "lz4-compression")]
 mod compression_lz4_block;
@@ -216,6 +216,72 @@ pub(crate) mod tests {
         )
     }
 
+    // Pathological schema with many numeric fields, as suggested in
+    // https://github.com/quickwit-oss/tantivy/issues/903 . This exercises the
+    // variable-length field-id and integer encoding heavily and verifies that
+    // documents still round-trip through the store correctly.
+    #[test]
+    fn test_store_many_numeric_fields_roundtrip() -> crate::Result<()> {
+        use common::DateTime;
+
+        use crate::schema::{OwnedValue, Value};
+
+        const NUM_U64_FIELDS: usize = 100;
+
+        let mut schema_builder = Schema::builder();
+        let u64_fields: Vec<_> = (0..NUM_U64_FIELDS)
+            .map(|i| schema_builder.add_u64_field(&format!("u64_{i}"), STORED))
+            .collect();
+        let i64_field = schema_builder.add_i64_field("i64", STORED);
+        let date_field = schema_builder.add_date_field("date", STORED);
+        let schema = schema_builder.build();
+
+        let path = Path::new("store");
+        let directory = RamDirectory::create();
+        let store_wrt = directory.open_write(path)?;
+        {
+            let mut store_writer =
+                StoreWriter::new(store_wrt, Compressor::default(), BLOCK_SIZE, true)?;
+            for i in 0..NUM_DOCS as u64 {
+                let mut doc = TantivyDocument::default();
+                for (j, field) in u64_fields.iter().enumerate() {
+                    doc.add_u64(*field, i.wrapping_mul(j as u64 + 1));
+                }
+                // Include both signs to exercise zig-zag encoding.
+                doc.add_i64(i64_field, (i as i64) - (NUM_DOCS as i64) / 2);
+                doc.add_date(date_field, DateTime::from_timestamp_nanos(i as i64 * 1_000));
+                store_writer.store(&doc, &schema)?;
+            }
+            store_writer.close()?;
+        }
+
+        let store_file = directory.open_read(path)?;
+        let store = StoreReader::open(store_file, 10)?;
+        for i in 0..NUM_DOCS as u64 {
+            let doc = store.get::<TantivyDocument>(i as u32)?;
+            for (j, field) in u64_fields.iter().enumerate() {
+                assert_eq!(
+                    doc.get_first(*field).unwrap().as_value().as_u64().unwrap(),
+                    i.wrapping_mul(j as u64 + 1)
+                );
+            }
+            assert_eq!(
+                doc.get_first(i64_field)
+                    .unwrap()
+                    .as_value()
+                    .as_i64()
+                    .unwrap(),
+                (i as i64) - (NUM_DOCS as i64) / 2
+            );
+            let expected_date = DateTime::from_timestamp_nanos(i as i64 * 1_000);
+            match doc.get_first(date_field).unwrap().as_value().into() {
+                OwnedValue::Date(dt) => assert_eq!(dt, expected_date),
+                other => panic!("expected date, got {other:?}"),
+            }
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_store_with_delete() -> crate::Result<()> {
         let mut schema_builder = schema::Schema::builder();

src/store/reader.rs

@@ -31,12 +31,17 @@ type Block = OwnedBytes;
 pub(crate) enum DocStoreVersion {
     V1 = 1,
     V2 = 2,
+    /// Same as V2 but field ids and integers are stored using variable-length
+    /// encoding (`VInt`, with zig-zag for signed integers) to make blocks more
+    /// compact before compression.
+    V3 = 3,
 }
 impl Display for DocStoreVersion {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             DocStoreVersion::V1 => write!(f, "V1"),
             DocStoreVersion::V2 => write!(f, "V2"),
+            DocStoreVersion::V3 => write!(f, "V3"),
         }
     }
 }
@@ -49,6 +54,7 @@ impl BinarySerializable for DocStoreVersion {
         Ok(match u32::deserialize(reader)? {
             1 => DocStoreVersion::V1,
             2 => DocStoreVersion::V2,
+            3 => DocStoreVersion::V3,
             v => {
                 return Err(io::Error::new(
                     io::ErrorKind::InvalidData,
@@ -499,7 +505,7 @@ mod tests {
         assert_eq!(store.cache_stats().cache_hits, 1);
         assert_eq!(store.cache_stats().cache_misses, 2);
 
-        assert_eq!(store.cache.peek_lru(), Some(232206));
+        assert_eq!(store.cache.peek_lru(), Some(229266));
 
         Ok(())
     }