From 4e97e38177190f865db51b8986b54795e55b934d Mon Sep 17 00:00:00 2001
From: ManyTheFish <many@meilisearch.com>
Date: Wed, 9 Oct 2024 14:39:27 +0200
Subject: [PATCH] Serialize docids bitmap one time

---
 milli/src/update/new/channel.rs | 14 +++++++++++---
 milli/src/update/new/merger.rs  | 11 +----------
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/milli/src/update/new/channel.rs b/milli/src/update/new/channel.rs
index bcac0fa03..8226046e6 100644
--- a/milli/src/update/new/channel.rs
+++ b/milli/src/update/new/channel.rs
@@ -5,6 +5,7 @@ use crossbeam_channel::{IntoIter, Receiver, SendError, Sender};
 use grenad::Merger;
 use heed::types::Bytes;
 use memmap2::Mmap;
+use roaring::RoaringBitmap;
 
 use super::extract::FacetKind;
 use super::StdResult;
@@ -46,6 +47,13 @@ impl KeyValueEntry {
         KeyValueEntry::SmallInMemory { key_length: key.len(), data: data.into_boxed_slice() }
     }
 
+    pub fn from_small_key_bitmap(key: &[u8], bitmap: RoaringBitmap) -> Self {
+        let mut data = Vec::with_capacity(key.len() + bitmap.serialized_size());
+        data.extend_from_slice(key);
+        bitmap.serialize_into(&mut data).unwrap();
+        KeyValueEntry::SmallInMemory { key_length: key.len(), data: data.into_boxed_slice() }
+    }
+
     pub fn from_large_key_value(key: &[u8], value: Mmap) -> Self {
         KeyValueEntry::LargeOnDisk { key: key.to_vec().into_boxed_slice(), value }
     }
@@ -232,10 +240,10 @@ impl MergerSender {
         DocumentsSender(self)
     }
 
-    pub fn send_documents_ids(&self, bitmap: &[u8]) -> StdResult<(), SendError<()>> {
-        let entry = EntryOperation::Write(KeyValueEntry::from_small_key_value(
+    pub fn send_documents_ids(&self, documents_ids: RoaringBitmap) -> StdResult<(), SendError<()>> {
+        let entry = EntryOperation::Write(KeyValueEntry::from_small_key_bitmap(
             DOCUMENTS_IDS_KEY.as_bytes(),
-            bitmap,
+            documents_ids,
         ));
         match self.send(WriterOperation { database: Database::Main, entry }) {
             Ok(()) => Ok(()),
diff --git a/milli/src/update/new/merger.rs b/milli/src/update/new/merger.rs
index 998a5d4a2..6183beb63 100644
--- a/milli/src/update/new/merger.rs
+++ b/milli/src/update/new/merger.rs
@@ -189,9 +189,7 @@ pub fn merge_grenad_entries(
         let _entered = span.enter();
 
         // Send the documents ids unionized with the current one
-        /// TODO return the slice of bytes directly
-        serialize_bitmap_into_vec(&documents_ids, &mut buffer);
-        sender.send_documents_ids(&buffer).unwrap();
+        sender.send_documents_ids(documents_ids).unwrap();
     }
 
     // ...
@@ -447,10 +445,3 @@ fn cbo_bitmap_serialize_into_vec<'b>(bitmap: &RoaringBitmap, buffer: &'b mut Vec
     CboRoaringBitmapCodec::serialize_into(bitmap, buffer);
     buffer.as_slice()
 }
-
-/// TODO Return the slice directly from the serialize_into method
-fn serialize_bitmap_into_vec(bitmap: &RoaringBitmap, buffer: &mut Vec<u8>) {
-    buffer.clear();
-    bitmap.serialize_into(buffer).unwrap();
-    // buffer.as_slice()
-}