[DistDGL][Optimizations]Rehash code to optimize for loop (#5224)

* Rehash code to optimize for loop

Reduced number of instructions in for loop, which exchanging edge features. This will reduce the number of times numpy's intersect1d is invoked (saving the runtime and memory overhead needs of numpy).

* Applying lintrunner patch to data_shuffle.py

tools/distpartitioning/data_shuffle.py

@@ -390,8 +390,8 @@ def exchange_feature(
 
     tokens = feat_key.split("/")
     assert len(tokens) == 3
+
     local_feat_key = "/".join(tokens[:-1]) + "/" + str(local_part_id)
-    for idx in range(world_size):
     # Get the partition ids for the range of global nids.
     if feat_type == constants.STR_NODE_FEATURES:
         # Retrieve the partition ids for the node features.
@@ -416,6 +416,7 @@ def exchange_feature(
         assert np.all(global_eids == data[constants.GLOBAL_EID][idx1])
         partid_slice = id_lookup.get_partition_ids(global_dst_nids)
 
+    for idx in range(world_size):
         cond = partid_slice == (idx + local_part_id * world_size)
         gids_per_partid = gids_feat[cond]
         tids_per_partid = tids_feat[cond]