Repository: meilisearch/meilisearch
File: crates/milli/src/update/new/words_prefix_docids.rs (Rust)

Code before:
                .bitmaps(prefix)
                .unwrap()
                .iter()
                .map(|bytes| CboRoaringBitmapCodec::deserialize_from(bytes))
                .union()?;

            buffer.clear();
            CboRoaringBitmapCodec::serialize_into_vec(&output, buffer);
            index.push(PrefixEntry { prefix, serialized_length: buffer.len() });
            file.write_all(buffer)
        })?;

        drop(frozen);

        // We iterate over all the collected and serialized bitmaps through
        // the files and entries to eventually put them in the final database.
        for refcell in local_entries {
            let (index, file, mut buffer) = refcell.into_inner();
            let mut file = file.into_inner().map_err(|e| e.into_error())?;
            file.rewind()?;
            let mut file = BufReader::new(file);
            for PrefixEntry { prefix, serialized_length } in index {
                buffer.resize(serialized_length, 0);
                file.read_exact(&mut buffer)?;
                self.prefix_database.remap_data_type::<Bytes>().put(
                    wtxn,
                    prefix.as_bytes(),
                    &buffer,
                )?;
            }
        }

        Ok(())
    }
}

/// Represents a prefix and the lenght the bitmap takes on disk.
struct PrefixEntry<'a> {
    prefix: &'a str,
    serialized_length: usize,
}

/// Stores prefixes along with all the pointers to the associated
/// CBoRoaringBitmaps.
///
/// They are collected synchronously and stored into an HashMap. The
/// Synchronous process is doing a small amount of work by just storing
/// pointers. It can then be accessed in parallel to get the associated
/// bitmaps pointers.
struct FrozenPrefixBitmaps<'a, 'rtxn> {
    prefixes_bitmaps: HashMap<&'a str, Vec<&'rtxn [u8]>>,

Review comment (performance):
It works for me.

However, prefixes have several sizes and can share the same ranges of words; for instance, the range of `hel` is a subset of the range of 'h'.
I don't know if it would improve the performance if we tried to iterate once on each word `O(words)` instead of iterating on each word for each prefix `O(words*3)`, 3 represents the maximum number of prefixes for a word, for example, `hello` can have `hel`, `he` and `h`.

Code after:
            .map(|(thread_id, rtxn)| {
                // `indexes` represent offsets at which prefixes computations were stored in the `file`.
                let mut indexes = Vec::new();
                let mut file = BufWriter::new(spooled_tempfile(
                    self.max_memory_by_thread.unwrap_or(usize::MAX),
                ));

                let mut buffer = Vec::new();
                for (prefix_index, prefix) in prefix_to_compute.iter().enumerate() {
                    // Is prefix for another thread?
                    if prefix_index % thread_count != thread_id {
                        continue;
                    }

                    let output = self
                        .database
                        .prefix_iter(&rtxn, prefix.as_bytes())?
                        .remap_types::<Str, CboRoaringBitmapCodec>()
                        .map(|result| result.map(|(_word, bitmap)| bitmap))
                        .union()?;

                    buffer.clear();
                    CboRoaringBitmapCodec::serialize_into_vec(&output, &mut buffer);
                    indexes.push(PrefixEntry { prefix, serialized_length: buffer.len() });
                    file.write_all(&buffer)?;
                }

                Ok((indexes, file))
            })
            .collect::<Result<Vec<_>>>()?;

        // We iterate over all the collected and serialized bitmaps through
        // the files and entries to eventually put them in the final database.
        let mut buffer = Vec::new();
        for (index, file) in outputs {
            let mut file = file.into_inner().map_err(|e| e.into_error())?;
            file.rewind()?;
            let mut file = BufReader::new(file);
            for PrefixEntry { prefix, serialized_length } in index {
                buffer.resize(serialized_length, 0);
                file.read_exact(&mut buffer)?;
                self.prefix_database.remap_data_type::<Bytes>().put(
                    wtxn,
                    prefix.as_bytes(),
                    &buffer,
                )?;
            }
        }

        Ok(())
    }