From 864f55a34ae7ecbf6c5884c80f2da43df7f4cb88 Mon Sep 17 00:00:00 2001 From: Shanicky Chen Date: Wed, 17 Jul 2024 21:32:34 +0800 Subject: [PATCH] Refactor ActorMapping iteration and worker lookup in mapping.rs --- src/common/src/hash/consistent_hash/mapping.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/common/src/hash/consistent_hash/mapping.rs b/src/common/src/hash/consistent_hash/mapping.rs index f95d89bad0cbd..d78d902fe7efb 100644 --- a/src/common/src/hash/consistent_hash/mapping.rs +++ b/src/common/src/hash/consistent_hash/mapping.rs @@ -332,7 +332,12 @@ impl ActorMapping { /// Transform the actor mapping to the worker slot mapping. Note that the parameter is a mapping from actor to worker. pub fn to_worker_slot(&self, actor_to_worker: &HashMap) -> WorkerSlotMapping { let mut worker_actors = HashMap::new(); - for (actor_id, worker_id) in actor_to_worker { + for actor_id in self.iter_unique() { + let worker_id = actor_to_worker + .get(&actor_id) + .cloned() + .unwrap_or_else(|| panic!("location for actor {} not found", actor_id)); + worker_actors .entry(worker_id) .or_insert(BTreeSet::new()) @@ -342,7 +347,7 @@ impl ActorMapping { let mut actor_location = HashMap::new(); for (worker, actors) in worker_actors { for (idx, &actor) in actors.iter().enumerate() { - actor_location.insert(*actor, WorkerSlotId::new(*worker, idx)); + actor_location.insert(actor, WorkerSlotId::new(worker, idx)); } }