Skip to content

Generic distributed query algorithm #3793

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jotare wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Generic distributed query algorithm #3793

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d2c7591
Move towards a generic distributed query algorithm
jotare Jul 9, 2026
ff2b3eb
Add macro to simplify gRPC function code
jotare Jul 9, 2026
3445466
Rename macro. Move merge back inside distributed_query
jotare Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 90 additions & 40 deletions nidx/src/searcher/grpc.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,31 +154,33 @@ macro_rules! shard_request {
}};
}

#[tonic::async_trait]
impl NidxSearcher for SearchServer {
async fn search(&self, request: Request<SearchRequest>) -> Result<Response<SearchResponse>> {
let coordinator = !request.metadata().contains_key(HEADER_LOCAL_ONLY);
let request = request.into_inner();
/// Operation to multiple shards distributed across the searcher cluster
macro_rules! shards_request {
($self:ident, $op:ty, $request:ident) => {{
let coordinator = !$request.metadata().contains_key(HEADER_LOCAL_ONLY);
let request = $request.into_inner();

let mut shards = vec![];
for shard_id in &request.shard_ids {
let shard_id = Uuid::parse_str(shard_id).map_err(NidxError::from)?;
shards.push(shard_id);
}
let shards = validate_shards(&request.shard_ids)?;

let response = if coordinator {
// We are the requested searcher to perform a distributed query
self.distributed_query(request, shards).await?
$self.distributed_query::<$op>(request, shards).await?
} else {
// This is a hopped node, i.e., another searcher has delegated this part of the query to
// us. We must query our shards and return either a full or partial response or an error.
// We won't hop to any other node
shard_search::search(shards, Arc::clone(&self.index_cache), request)
<$op>::local_query(Arc::clone(&$self.index_cache), request, shards)
.await?
.data
};

Ok(Response::new(response))
}};
}

#[tonic::async_trait]
impl NidxSearcher for SearchServer {
async fn search(&self, request: Request<SearchRequest>) -> Result<Response<SearchResponse>> {
shards_request!(self, SearchOp, request)
}

async fn suggest(&self, request: Request<SuggestRequest>) -> Result<Response<SuggestResponse>> {
Expand Down Expand Up @@ -271,7 +273,10 @@ impl SearchServer {
/// requested to query several partitions and one or more can fail, it will then return a
/// partial result and is our responsibility to retry on other nodes.
///
async fn distributed_query(&self, request: SearchRequest, shards: Vec<Uuid>) -> NidxResult<SearchResponse> {
async fn distributed_query<Op>(&self, request: Op::Request, shards: Vec<Uuid>) -> NidxResult<Op::Response>
where
Op: SearcherOp,
{
let mut partitioner = QueryPartitioner::new(&self.shard_selector, &shards)?;

// Use the query partitioner to group requests by the most preferred
Expand Down Expand Up @@ -310,7 +315,7 @@ impl SearchServer {
let mut tasks = JoinSet::new();
for (node, shard_ids) in groups {
tasks.spawn(
self.node_query(node, shard_ids, request.clone())
self.node_query::<Op>(node, shard_ids, request.clone())
.await?
.instrument(Span::current()),
);
Expand Down Expand Up @@ -353,45 +358,36 @@ impl SearchServer {
}
}

let merged = if responses.len() == 1 {
responses.pop().unwrap()
} else {
shard_merge::merge(
responses,
OrderBy::from(&request),
Limit(request.result_per_page as usize),
)
};
let merged = Op::merge(&request, responses);
Ok(merged)
}

async fn node_query(
async fn node_query<Op>(
&self,
node: SearcherNode,
shard_ids: Vec<Uuid>,
request: SearchRequest,
) -> NidxResult<Pin<Box<dyn Future<Output = NidxResult<PartialResponse<SearchResponse>>> + Send + 'static>>> {
shards: Vec<Uuid>,
request: Op::Request,
) -> NidxResult<Pin<Box<dyn Future<Output = NidxResult<PartialResponse<Op::Response>>> + Send + 'static>>>
where
Op: SearcherOp,
{
match node {
SearcherNode::This => {
let index_cache = Arc::clone(&self.index_cache);
Ok(Box::pin(async move {
shard_search::search(shard_ids, index_cache, request).await
}))
Ok(Box::pin(
async move { Op::local_query(index_cache, request, shards).await },
))
}
SearcherNode::Remote(ref hostname) => {
match self.get_client(hostname).await {
Ok(client) => Ok(Box::pin(async move {
Self::remote_query(client, request, shard_ids)
.await
.map_err(NidxError::GrpcError)
})),
Ok(client) => Ok(Box::pin(async move { Op::remote_query(client, request, shards).await })),
Err(e) => {
// A client for this node is not available, although we've seen it
// in the k8s cluster. We skip it and will retry those shards on
// another node to complete the request
warn!(
?node,
?shard_ids,
?shards,
"{}",
format!("Error getting a client for node '{hostname}': {e}")
);
Expand All @@ -401,12 +397,47 @@ impl SearchServer {
}
}
}
}

trait SearcherOp: Clone + Send {
type Request: Clone + Send + 'static;
type Response: Send + 'static;

fn local_query(
index_cache: Arc<IndexCache>,
request: Self::Request,
shards: Vec<Uuid>,
) -> impl std::future::Future<Output = NidxResult<PartialResponse<Self::Response>>> + Send;

fn remote_query(
client: SearcherClient,
request: Self::Request,
shards: Vec<Uuid>,
) -> impl std::future::Future<Output = NidxResult<PartialResponse<Self::Response>>> + Send;

fn merge(request: &Self::Request, partitions: Vec<Self::Response>) -> Self::Response;
}

#[derive(Clone)]
struct SearchOp;

impl SearcherOp for SearchOp {
type Request = SearchRequest;
type Response = SearchResponse;

async fn local_query(
index_cache: Arc<IndexCache>,
request: Self::Request,
shards: Vec<Uuid>,
) -> NidxResult<PartialResponse<Self::Response>> {
shard_search::search(shards, index_cache, request).await
}

async fn remote_query(
mut client: SearcherClient,
request: SearchRequest,
request: Self::Request,
shards: Vec<Uuid>,
) -> tonic::Result<PartialResponse<SearchResponse>> {
) -> NidxResult<PartialResponse<Self::Response>> {
// Send the query to a different node specifying only a subset of shards
let mut search_request = request.clone();
search_request.shard_ids = shards.iter().map(|x| x.to_string()).collect();
Expand All @@ -416,7 +447,7 @@ impl SearchServer {
let mut request = Request::new(search_request);
request.metadata_mut().insert(HEADER_LOCAL_ONLY, "1".parse().unwrap());

let response = client.search(request).await?.into_inner();
let response = client.search(request).await.map_err(NidxError::GrpcError)?.into_inner();

let result = if response.shard_ids.len() == shards.len() {
PartialResponse { data: response, shards }
Expand All @@ -433,6 +464,16 @@ impl SearchServer {
};
Ok(result)
}

fn merge(request: &Self::Request, mut partitions: Vec<Self::Response>) -> Self::Response {
if partitions.len() == 1 {
partitions.pop().unwrap()
} else {
let order_by = OrderBy::from(request);
let limit = Limit(request.result_per_page as usize);
shard_merge::merge(partitions, order_by, limit)
}
}
}

struct QueryPartitioner {
Expand Down Expand Up @@ -462,3 +503,12 @@ impl QueryPartitioner {
if !nodes.is_empty() { Some(nodes.remove(0)) } else { None }
}
}

fn validate_shards(shards: &[String]) -> NidxResult<Vec<Uuid>> {
let mut valid = Vec::with_capacity(shards.len());
for shard_id in shards {
let shard_id = Uuid::parse_str(shard_id).map_err(NidxError::from)?;
valid.push(shard_id);
}
Ok(valid)
}
Loading