Function meta_embed_search

nodedb-vector/src/multivec/meta_embed.rs:34–74

Search a `MultiVectorStore` using budgeted MaxSim with optional PLAID candidate pruning. # Parameters `store` — the document collection. `plaid` — optional PLAID pruner (pass `None` to scan all docs). `query` — query Meta Token vectors (Matryoshka ordering). `budget` — number of leading query tokens to use; 0 falls back to all. `k` — number of top documents to return. `metric` — distance

(
    store: &MultiVectorStore,
    plaid: Option<&PlaidPruner>,
    query: &[Vec<f32>],
    budget: u8,
    k: usize,
    metric: DistanceMetric,
)

/// # Returns
/// A `Vec<(doc_id, score)>` sorted descending by score, length ≤ `k`.
pub fn meta_embed_search(
    store: &MultiVectorStore,
    plaid: Option<&PlaidPruner>,
    query: &[Vec<f32>],
    budget: u8,
    k: usize,
    metric: DistanceMetric,
) -> Vec<(u32, f32)> {
    if k == 0 || query.is_empty() {
        return Vec::new();
    }

    // Effective budget: 0 means use all query vectors.
    let effective_budget = if budget == 0 {
        query.len() as u8
    } else {
        budget
    };

    // Determine candidate set.
    let candidate_ids: Vec<u32> = match plaid {
        Some(pruner) => pruner.candidates(query),
        None => store.iter().map(|doc| doc.doc_id).collect(),
    };

    // Score each candidate.
    let mut scored: Vec<(u32, f32)> = candidate_ids
        .into_iter()
        .filter_map(|doc_id| {
            store.get(doc_id).map(|doc| {
                let score = budgeted_maxsim(query, &doc.vectors, effective_budget, metric);
                (doc_id, score)
            })
        })
        .collect();

    // Sort descending by score.
    scored.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    scored.truncate(k);
    scored
}

// ---------------------------------------------------------------------------
// Tests