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

Code before:
            .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) {
                        keys_to_delete.push(key_bytes.to_vec().into_boxed_slice());
                    }
                }

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

    Ok(results)
}

/// Rebuilds the hannoy graph and do not touch to the embeddings.
///
/// This follows a bug in hannoy v0.0.9 and v0.1.0 where the graph
/// was not built correctly.
pub(super) struct RebuildHannoyGraph();

impl UpgradeIndex for RebuildHannoyGraph {
    fn upgrade(
        &self,
        wtxn: &mut RwTxn,
        index: &Index,
        must_stop_processing: &MustStopProcessing,
        progress: Progress,
    ) -> Result<bool> {
        let embedders = index.embedding_configs();
        let backend = index.get_vector_store(wtxn)?.unwrap_or_default();

        for config in embedders.embedding_configs(wtxn)? {
            let embedder_info = embedders.embedder_info(wtxn, &config.name)?.unwrap();
            let mut vector_store = VectorStore::new(

Review comment (performance):
Do we really need to do this? That's a lot of allocations for something that is poised to deletion anyway. why not just sending batch of keys to be deleted in a different thread instead? Like using a channel to communicate with the other thread?

Code after:
/// 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
    // selecting ranges of keys that will be processed by each thread. We
    // also make sure not to specify the first and last keys to make sure
    // that if the fields to clean have keys that are higher or lower than
    // the first or last keys in the word dictionary we still find them.

    // 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().into_boxed_slice()));
        }

        count += 1;
    }

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

    Ok(bounds)
}

/// Rebuilds the hannoy graph and do not touch to the embeddings.