Repository: meilisearch/meilisearch
File: crates/milli/src/update/upgrade/v1_32.rs (Rust)

Code before:

    // Here We make sure to start with an unbounded
    // left bound for the first range
    let mut bounds = vec![None];
    let mut stream = fst.stream();
    let mut count = 0;
    while let Some(key) = stream.next() {
        let is_first = count == 0;
        let is_last = count == fst.len() - 1;

        // In this loop we make sure to account for every bounds
        // to divide the work between threads but not send the bounds
        // for the beginning or the end of the word dictionary
        if count % keys_by_thread == 0 && !(is_first || is_last) {
            bounds.push(Some(key.to_vec()));
        }

        count += 1;
    }

    // We now push the last bound that
    // defines the end of the last range
    bounds.push(None);

    // We create a thread pool and generate enough read transactions for each one of them.
    let pool = rayon::ThreadPoolBuilder::new().num_threads(threads_count).build()?;
    let rtxns = iter::repeat_with(|| index.env.nested_read_txn(wtxn))
        .take(bounds.len().saturating_sub(1))
        .collect::<heed::Result<Vec<_>>>()?;

    let results = pool.install(|| {
        rtxns
            .into_par_iter()
            .zip_eq(bounds.windows(2).collect::<Vec<_>>())
            .map(|(rtxn, win)| {
                let bound = match [win[0].as_deref(), win[1].as_deref()] {
                    [None, None] => (Bound::Unbounded, Bound::Unbounded),
                    [None, Some(end)] => (Bound::Unbounded, Bound::Excluded(end)),
                    [Some(start), None] => (Bound::Included(start), Bound::Unbounded),
                    [Some(start), Some(end)] => (Bound::Included(start), Bound::Excluded(end)),
                };

                let mut keys_to_delete = Vec::new();
                let iter = database.remap_types::<Bytes, DecodeIgnore>().range(&rtxn, &bound);
                for result in iter? {
                    let (key_bytes, ()) = result?;
                    let (_word, fid) =
                        StrBEU16Codec::bytes_decode(key_bytes).map_err(heed::Error::Decoding)?;

                    // TODO optimize allocations
                    if fids_to_delete.contains(&fid) {

Review comment (suggestion):
We should create that pool only once, maybe in the `delete_old_fid_based_databases_from_fids` so we don't waste time re-creating a threadpool. `delete_old_fid_word_count_docids` is called 3 times, so we create 3 threadpools in total.

Code after:
    let results = pool.install(|| {
        rtxns
            .into_par_iter()
            .zip_eq(bounds.windows(2).collect::<Vec<_>>())
            .map(|(rtxn, win)| {
                let bound = match [win[0].as_deref(), win[1].as_deref()] {
                    [None, None] => (Bound::Unbounded, Bound::Unbounded),
                    [None, Some(end)] => (Bound::Unbounded, Bound::Excluded(end)),
                    [Some(start), None] => (Bound::Included(start), Bound::Unbounded),
                    [Some(start), Some(end)] => (Bound::Included(start), Bound::Excluded(end)),
                };

                let mut keys_to_delete = Vec::new();
                let iter = database.remap_types::<Bytes, DecodeIgnore>().range(&rtxn, &bound);
                for result in iter? {
                    let (key_bytes, ()) = result?;
                    let (_word, fid) =
                        StrBEU16Codec::bytes_decode(key_bytes).map_err(heed::Error::Decoding)?;

                    if fids_to_delete.contains(&fid) {
                        keys_to_delete.push(key_bytes.to_vec().into_boxed_slice());
                    }
                }

                Ok(keys_to_delete) as crate::Result<_>
            })
            .collect_vec_list()
    });

    Ok(results)
}

/// Use the FST to balance the work between threads
/// by generating appropriate word bounds.
fn compute_bounds<'txn>(
    wtxn: &mut RwTxn<'txn>,
    index: &Index,
) -> crate::Result<Vec<Option<Box<[u8]>>>> {
    let fst = index.words_fst(wtxn)?;

    let threads_count = rayon::current_num_threads() * 4;
    let keys_by_thread = fst.len().div_ceil(threads_count);

    // We iterate over the FST keys that represents the word dictionary and
    // roughly represents what can be found in the database we are cleaning.
    //
    // The database we are cleaning contains different words from the word
    // dictionary as it contains words from fields that are not indexed too
    // but it is mixed with indexed ones.
    //
    // We then divide equally the entries of the database to clean by