From 7905d21ff4640dfca4dfd4c1694879d8cc8870c7 Mon Sep 17 00:00:00 2001 From: Shanicky Chen Date: Mon, 1 Jul 2024 00:03:21 +0800 Subject: [PATCH] Refactor rebuild_fragment_mapping, rename var and change storage to nested HashMaps in utils.rs --- src/meta/src/controller/utils.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/meta/src/controller/utils.rs b/src/meta/src/controller/utils.rs index 876dcfb3037d5..a3855903e45aa 100644 --- a/src/meta/src/controller/utils.rs +++ b/src/meta/src/controller/utils.rs @@ -879,13 +879,16 @@ pub fn rebuild_fragment_mapping_from_actors( WorkerId, )>, ) -> Vec { - let mut actor_locations = HashMap::new(); + let mut all_actor_locations = HashMap::new(); let mut actor_bitmaps = HashMap::new(); let mut fragment_actors = HashMap::new(); let mut fragment_dist = HashMap::new(); for (fragment_id, dist, actor_id, bitmap, worker_id) in job_actors { - actor_locations.insert(actor_id as hash::ActorId, worker_id as u32); + all_actor_locations + .entry(fragment_id) + .or_insert(HashMap::new()) + .insert(actor_id as hash::ActorId, worker_id as u32); actor_bitmaps.insert(actor_id, bitmap); fragment_actors .entry(fragment_id) @@ -896,6 +899,7 @@ pub fn rebuild_fragment_mapping_from_actors( let mut result = vec![]; for (fragment_id, dist) in fragment_dist { + let mut actor_locations = all_actor_locations.remove(&fragment_id).unwrap(); let fragment_worker_slot_mapping = match dist { DistributionType::Single => { let actor = fragment_actors