Function finalize_ann_results

src/vector_store.rs:404–445

(
    collection: &Collection,
    hnsw_index: &HNSWIndex,
    results: Vec<(SharedLatestNode, MetricResult)>,
    query: &[f32],
    top_k: Option<usize>,
    return_raw_text: bool,
)

}

pub fn finalize_ann_results(
    collection: &Collection,
    hnsw_index: &HNSWIndex,
    results: Vec<(SharedLatestNode, MetricResult)>,
    query: &[f32],
    top_k: Option<usize>,
    return_raw_text: bool,
) -> Result<Vec<InternalSearchResult>, WaCustomError> {
    let filtered = remove_duplicates_and_filter(hnsw_index, results, top_k, &hnsw_index.cache);
    let mut results = Vec::with_capacity(top_k.unwrap_or(filtered.len()));
    let mag_query = query.iter().map(|x| x * x).sum::<f32>().sqrt();

    for (internal_id, _) in filtered {
        let raw_emb = collection
            .get_raw_emb_by_internal_id(&internal_id)
            .ok_or_else(|| {
                WaCustomError::NotFound(format!("raw embedding not found for id={internal_id:?}"))
            })?;
        let dense_values = raw_emb.dense_values.as_ref().ok_or_else(|| {
            WaCustomError::NotFound("dense values not found for raw embedding".to_string())
        })?;
        let dp = dot_product_f32(query, dense_values);
        let mag_raw = dense_values.iter().map(|x| x * x).sum::<f32>().sqrt();
        let cs = dp / (mag_query * mag_raw);
        results.push((
            internal_id,
            Some(raw_emb.id.clone()),
            raw_emb.document_id.clone(),
            cs,
            if return_raw_text {
                raw_emb.text.clone()
            } else {
                None
            },
        ));
    }
    results.sort_unstable_by(|(_, _, _, a, _), (_, _, _, b, _)| b.total_cmp(a));
    if let Some(k) = top_k {
        results.truncate(k);
    }
    Ok(results)
}

/// Intermediate representation of the embedding in a form that's
/// ready for indexing.