From 519905ef9c3c20a759555395ace1d04380e27c71 Mon Sep 17 00:00:00 2001 From: ManyTheFish Date: Mon, 6 Oct 2025 12:10:48 +0200 Subject: [PATCH] Fix remote index collision with HashMap-based lookup - Replace BTreeMap with HashMap for (remote, index_uid) -> primary_key lookup - Prevents collisions when multiple remotes have same index_uid but different primary keys --- .../src/search/federated/perform.rs | 50 +++++++++++-------- 1 file changed, 29 insertions(+), 21 deletions(-) diff --git a/crates/meilisearch/src/search/federated/perform.rs b/crates/meilisearch/src/search/federated/perform.rs index e7f9bf4c8..3db6fb43d 100644 --- a/crates/meilisearch/src/search/federated/perform.rs +++ b/crates/meilisearch/src/search/federated/perform.rs @@ -132,14 +132,18 @@ pub async fn perform_federated_search( let query_metadata = if include_metadata { let mut query_metadata = Vec::new(); - // Create a map of remote results by (index_uid, remote) for quick lookup + // Create a map of (remote, index_uid) -> primary_key for quick lookup // This prevents collisions when multiple remotes have the same index_uid but different primary keys - let mut remote_results_by_index_and_remote = std::collections::BTreeMap::new(); + let mut remote_primary_keys = std::collections::HashMap::new(); for remote_result in &remote_results { if let Some(remote_metadata) = &remote_result.metadata { for remote_meta in remote_metadata { - let key = (remote_meta.index_uid.clone(), remote_meta.remote.clone()); - remote_results_by_index_and_remote.insert(key, remote_meta.clone()); + if let Some(remote_name) = &remote_meta.remote { + let key = (remote_name.clone(), remote_meta.index_uid.clone()); + if let Some(primary_key) = &remote_meta.primary_key { + remote_primary_keys.insert(key, primary_key.clone()); + } + } } } } @@ -154,25 +158,22 @@ pub async fn perform_federated_search( let remote = federation_options.and_then(|options| options.remote); // Get primary key for this index - let mut primary_key = None; - - if remote.is_some() { - // For remote queries, try to get primary key from remote results - // Use composite key (index_uid, remote) to avoid collisions - let lookup_key = (index_uid.clone(), remote.clone()); - if let Some(remote_meta) = remote_results_by_index_and_remote.get(&lookup_key) { - primary_key = remote_meta.primary_key.clone(); + let primary_key = match &remote { + Some(remote_name) => { + // For remote queries, try to get primary key from remote results + // Use composite key (remote, index_uid) to avoid collisions + let lookup_key = (remote_name.clone(), index_uid.clone()); + remote_primary_keys.get(&lookup_key).cloned() } - } else { - // For local queries, get primary key from local index - primary_key = index_scheduler.index(&index_uid).ok().and_then(|index| { - index.read_txn().ok().and_then(|rtxn| { - let pk = index.primary_key(&rtxn).ok().flatten().map(|pk| pk.to_string()); - drop(rtxn); - pk + None => { + // For local queries, get primary key from local index + index_scheduler.index(&index_uid).ok().and_then(|index| { + index.read_txn().ok().and_then(|rtxn| { + index.primary_key(&rtxn).ok().flatten().map(|pk| pk.to_string()) + }) }) - }); - } + } + }; query_metadata.push(SearchMetadata { query_uid, index_uid, primary_key, remote }); } @@ -713,6 +714,13 @@ impl RemoteSearch { continue 'remote_queries; } + // Add remote name to metadata + if let Some(metadata) = res.metadata.as_mut() { + for meta in metadata { + meta.remote = Some(node_name.clone()); + } + } + federation.insert( FEDERATION_REMOTE.to_string(), serde_json::Value::String(node_name.clone()),