Repository: pyg-team/pytorch_geometric
Pull request #8083 (2023-09-27 14:50:18)

Update `GraphStore` and `FeatureStore` [1/6]
This code belongs to the part of the whole distributed training for PyG.

Please be aware that this PR should be merged before Loaders package! - @JakubPietrakIntel 
Loaders:
1.  https://github.com/pyg-team/pytorch_geometric/pull/8079
2.  https://github.com/pyg-team/pytorch_geometric/pull/8080
3.  https://github.com/pyg-team/pytorch_geometric/pull/8085

Other PRs related to this module:
DistSampler: https://github.com/pyg-team/pytorch_geometric/pull/7974

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8083?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8083](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8083?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (81206ee) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/4a7658f22cc16dfc4fc776195fb91d4257507c34?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4a7658f) will **decrease** coverage by `0.75%`.
> Report is 1 commits behind head on master.
> The diff coverage is `65.21%`.

```diff
@@            Coverage Diff             @@
##           master    #8083      +/-   ##
==========================================
- Coverage   89.42%   88.68%   -0.75%     
==========================================
  Files         467      467              
  Lines       27882    27902      +20     
==========================================
- Hits        24934    24745     -189     
- Misses       2948     3157     +209     
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8083?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/distributed/local\_graph\_store.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8083?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2xvY2FsX2dyYXBoX3N0b3JlLnB5) | `91.20% <ø> (ø)` | |
| [torch\_geometric/distributed/local\_feature\_store.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8083?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2xvY2FsX2ZlYXR1cmVfc3RvcmUucHk=) | `78.00% <12.50%> (-3.11%)` | :arrow_down: |
| [torch\_geometric/data/graph\_store.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8083?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZ3JhcGhfc3RvcmUucHk=) | `92.75% <76.31%> (-5.61%)` | :arrow_down: |

... and [31 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8083/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

JakubPietrakIntel: @rusty1s @ZhengHongming888 
I added [ac49add](https://github.com/pyg-team/pytorch_geometric/pull/8083/commits/ac49add09e2fe95b9771b5f9a4a97527d22d9ce8) which enforces sorting by destination node in LGS `edge_index` (in other words sort by col, or `edge_index[1]`) to match CSC format used for sampling. This is condition ensures that input data won't be permuted again at neighbor sampler init.
We do that to make sure that correct edge_ids are returned at sampler's output and in `filter_fn` in `NodeLoader` i.e. [L161](https://github.com/pyg-team/pytorch_geometric/blob/ac49add09e2fe95b9771b5f9a4a97527d22d9ce8/torch_geometric/loader/node_loader.py#L161) the expression is evaluated correctly with `perm = None`.
We have considered alternative solutions and decided that this one will be the best regarding performance, at a cost of adding a constraint on incoming data. We can't have a global perm stored in LGS, as METIS is not performing a simple split, so the ordering would be lost. The other solution would be to use local values per partition in `edge_id`, hence increasing complexity and adding yet another re-indexing step. Furthermore, it would require RPC requests for perm values located in remote partitions, because the sampler may return edge_ids located in remote partition.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6d1bf8db8..e5ce63214 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Update `GraphStore` and `FeatureStore` to support distributed training ([#8083](https://github.com/pyg-team/pytorch_geometric/pull/8083))
 - Disallow the usage of `add_self_loops=True` in `GCNConv(normalize=False)` ([#8210](https://github.com/pyg-team/pytorch_geometric/pull/8210))
 - Disable device asserts during `torch_geometric.compile` ([#8220](https://github.com/pyg-team/pytorch_geometric/pull/8220))
 
diff --git a/torch_geometric/data/graph_store.py b/torch_geometric/data/graph_store.py
index cd591895d..2985026e1 100644
--- a/torch_geometric/data/graph_store.py
+++ b/torch_geometric/data/graph_store.py
@@ -278,7 +278,16 @@ class GraphStore:
                 row = ptr2index(row)
 
             if attr.layout != EdgeLayout.CSC:  # COO->CSC
-                num_cols = attr.size[1] if attr.size else int(col.max()) + 1
+                if hasattr(self, 'meta') and self.meta.get('is_hetero', False):
+                    # Hotfix for `LocalGraphStore`, where in heterogeneous
+                    # graphs, edge indices for different edge types have
+                    # continuous indices not starting at 0.
+                    num_cols = int(col.max()) + 1
+                elif attr.size is not None:
+                    num_cols = attr.size[1]
+                else:
+                    num_cols = int(col.max()) + 1
+
                 if not attr.is_sorted:  # Not sorted by destination.
                     col, perm = index_sort(col, max_value=num_cols)
                     row = row[perm]
@@ -300,6 +309,17 @@ class GraphStore:
         store: bool = False,
     ) -> ConversionOutputType:
 
+        is_hetero = True  # Default.
+        if hasattr(self, 'meta'):  # `LocalGraphStore` hack.
+            is_hetero = self.meta.get('is_hetero', False)
+
+        if not is_hetero:
+            edge_attrs: List[EdgeAttr] = []
+            for attr in self.get_all_edge_attrs():
+                edge_attrs.append(attr)
+
+            return self._edge_to_layout(edge_attrs[0], layout, store)
+
         # Obtain all edge attributes, grouped by type:
         edge_type_attrs: Dict[EdgeType, List[EdgeAttr]] = defaultdict(list)
         for attr in self.get_all_edge_attrs():
diff --git a/torch_geometric/distributed/dist_neighbor_sampler.py b/torch_geometric/distributed/dist_neighbor_sampler.py
index 7031e6660..84ae7dcfb 100644
--- a/torch_geometric/distributed/dist_neighbor_sampler.py
+++ b/torch_geometric/distributed/dist_neighbor_sampler.py
@@ -100,7 +100,6 @@ class DistNeighborSampler:
         self.temporal_strategy = temporal_strategy
         self.time_attr = time_attr
         self.with_edge_attr = self.dist_feature.has_edge_attr()
-        _, _, self.edge_permutation = self.dist_graph.csc()
         self.csc = True
 
     def register_sampler_rpc(self) -> None:
@@ -121,6 +120,8 @@ class DistNeighborSampler:
             temporal_strategy=self.temporal_strategy,
             time_attr=self.time_attr,
         )
+        self.edge_permutation = self._sampler.perm
+
         rpc_sample_callee = RPCSamplingCallee(self._sampler)
         self.rpc_sample_callee_id = rpc_register(rpc_sample_callee)
 
@@ -660,9 +661,6 @@ class DistNeighborSampler:
                 efeats = None
 
         output.metadata = (*output.metadata, nfeats, nlabels, efeats)
-        if self.is_hetero:
-            output.row = remap_keys(output.row, self._sampler.to_edge_type)
-            output.col = remap_keys(output.col, self._sampler.to_edge_type)
         return output
 
     def __repr__(self) -> str:
diff --git a/torch_geometric/distributed/local_feature_store.py b/torch_geometric/distributed/local_feature_store.py
index ec86239e3..f22f7bb4c 100644
--- a/torch_geometric/distributed/local_feature_store.py
+++ b/torch_geometric/distributed/local_feature_store.py
@@ -49,20 +49,21 @@ class LocalFeatureStore(FeatureStore):
     """
     def __init__(self):
         super().__init__(tensor_attr_cls=LocalTensorAttr)
-
         self._feat: Dict[Tuple[Union[NodeType, EdgeType], str], Tensor] = {}
-
         # Save the global node/edge IDs:
         self._global_id: Dict[Union[NodeType, EdgeType], Tensor] = {}
-
         # Save the mapping from global node/edge IDs to indices in `_feat`:
         self._global_id_to_index: Dict[Union[NodeType, EdgeType], Tensor] = {}
-
-        # For partition/rpc information related to distribute features:
-        self.num_partitions = 1
-        self.partition_idx = 0
-        self.feature_pb: Union[Tensor, Dict[NodeOrEdgeType, Tensor]]
-        self.local_only = False
+        # For partition/RPC info related to distributed features:
+        self.num_partitions: int = 1
+        self.partition_idx: int = 0
+        # Mapping between node ID and partition ID:
+        self.node_feat_pb: Union[Tensor, Dict[NodeType, Tensor]]
+        # Mapping between edge ID and partition ID:
+        self.edge_feat_pb: Union[Tensor, Dict[EdgeType, Tensor]]
+        self.labels: Optional[Tensor] = None  # Node labels.
+
+        self.local_only: bool = False
         self.rpc_router: Optional[RPCRouter] = None
         self.meta: Optional[Dict] = None
         self.rpc_call_id: Optional[int] = None
@@ -147,6 +148,16 @@ class LocalFeatureStore(FeatureStore):
         else:
             self.rpc_call_id = None
 
+    def has_edge_attr(self) -> bool:
+        has_edge_attr = False
+        for k in [key for key in self._feat.keys() if 'edge_attr' in key]:
+            try:
+                self.get_tensor(k[0], 'edge_attr')
+                has_edge_attr = True
+            except KeyError:
+                pass
+        return has_edge_attr
+
     def lookup_features(
         self,
         index: Tensor,
@@ -164,8 +175,11 @@ class LocalFeatureStore(FeatureStore):
             try:
                 remote_feature_list = remote_fut.wait()
                 # combine the feature from remote and local
-                result = torch.zeros(index.size(0), local_feature[0].size(1),
-                                     dtype=local_feature[0].dtype)
+                result = torch.zeros(
+                    index.size(0),
+                    local_feature[0].size(1),
+                    dtype=local_feature[0].dtype,
+                )
                 result[local_feature[1]] = local_feature[0]
                 for remote in remote_feature_list:
                     result[remote[1]] = remote[0]
@@ -184,12 +198,7 @@ class LocalFeatureStore(FeatureStore):
         input_type: Optional[Union[NodeType, EdgeType]] = None,
     ) -> Tuple[Tensor, Tensor]:
         r"""Lookup the features in local nodes based on node/edge IDs."""
-        if self.meta['is_hetero']:
-            feat = self
-            pb = self.feature_pb[input_type]
-        else:
-            feat = self
-            pb = self.feature_pb
+        pb = self.node_feat_pb if is_node_feat else self.edge_feat_pb
 
         input_order = torch.arange(index.size(0), dtype=torch.long)
         partition_ids = pb[index]
@@ -198,23 +207,23 @@ class LocalFeatureStore(FeatureStore):
         local_ids = torch.masked_select(index, local_mask)
         local_index = torch.masked_select(input_order, local_mask)
 
-        if self.meta["is_hetero"]:
+        if self.meta['is_hetero']:
             if is_node_feat:
                 kwargs = dict(group_name=input_type, attr_name='x')
-                ret_feat = feat.get_tensor_from_global_id(
+                ret_feat = self.get_tensor_from_global_id(
                     index=local_ids, **kwargs)
             else:
                 kwargs = dict(group_name=input_type, attr_name='edge_attr')
-                ret_feat = feat.get_tensor_from_global_id(
+                ret_feat = self.get_tensor_from_global_id(
                     index=local_ids, **kwargs)
         else:
             if is_node_feat:
                 kwargs = dict(group_name=None, attr_name='x')
-                ret_feat = feat.get_tensor_from_global_id(
+                ret_feat = self.get_tensor_from_global_id(
                     index=local_ids, **kwargs)
             else:
                 kwargs = dict(group_name=(None, None), attr_name='edge_attr')
-                ret_feat = feat.get_tensor_from_global_id(
+                ret_feat = self.get_tensor_from_global_id(
                     index=local_ids, **kwargs)
 
         return ret_feat, local_index
@@ -226,10 +235,7 @@ class LocalFeatureStore(FeatureStore):
         input_type: Optional[Union[NodeType, EdgeType]] = None,
     ) -> torch.futures.Future:
         r"""Fetch the remote features with the remote node/edge IDs."""
-        if self.meta["is_hetero"]:
-            pb = self.feature_pb[input_type]
-        else:
-            pb = self.feature_pb
+        pb = self.node_feat_pb if is_node_feat else self.edge_feat_pb
 
         input_order = torch.arange(index.size(0), dtype=torch.long)
         partition_ids = pb[index]
@@ -237,7 +243,7 @@ class LocalFeatureStore(FeatureStore):
         for pidx in range(0, self.num_partitions):
             if pidx == self.partition_idx:
                 continue
-            remote_mask = (partition_ids == pidx)
+            remote_mask = partition_ids == pidx
             remote_ids = index[remote_mask]
             if remote_ids.shape[0] > 0:
                 to_worker = self.rpc_router.get_to_worker(pidx)
diff --git a/torch_geometric/distributed/local_graph_store.py b/torch_geometric/distributed/local_graph_store.py
index 8e6878cdd..01e215c2c 100644
--- a/torch_geometric/distributed/local_graph_store.py
+++ b/torch_geometric/distributed/local_graph_store.py
@@ -7,6 +7,7 @@ from torch import Tensor
 
 from torch_geometric.data import EdgeAttr, GraphStore
 from torch_geometric.typing import EdgeTensorType, EdgeType, NodeType
+from torch_geometric.utils import sort_edge_index
 
 
 class LocalGraphStore(GraphStore):
@@ -27,8 +28,8 @@ class LocalGraphStore(GraphStore):
         self.edge_pb: Union[Tensor, Dict[EdgeType, Tensor]] = None
         # Meta information related to partition and graph store info
         self.meta: Optional[Dict[Any, Any]] = None
-        # Partition labels
-        self.labels: Union[Tensor, Dict[EdgeType, Tensor]] = None
+        # If data is sorted based on destination nodes (CSC format):
+        self.is_sorted: Optional[bool] = None
 
     @staticmethod
     def key(attr: EdgeAttr) -> Tuple:
@@ -40,17 +41,11 @@ class LocalGraphStore(GraphStore):
         node_type: Optional[NodeType] = None,
     ) -> Tensor:
         r"""Get the partition IDs of node IDs for a specific node type."""
-        if self.meta['is_hetero']:
-            assert node_type is not None
-            return self.node_pb[node_type][ids]
         return self.node_pb[ids]
 
     def get_partition_ids_from_eids(self, eids: torch.Tensor,
                                     edge_type: Optional[EdgeType] = None):
         r"""Get the partition IDs of edge IDs for a specific edge type."""
-        if self.meta["is_hetero"]:
-            assert edge_type is not None
-            return self.edge_pb[edge_type][eids]
         return self.edge_pb[eids]
 
     def put_edge_id(self, edge_id: Tensor, *args, **kwargs) -> bool:
@@ -90,23 +85,37 @@ class LocalGraphStore(GraphStore):
         edge_id: Tensor,
         edge_index: Tensor,
         num_nodes: int,
+        is_sorted: bool = False,
     ) -> 'LocalGraphStore':
-        r"""Creates a local graph store from a homogeneous :pyg:`PyG` graph.
+        r"""Creates a local graph store from a homogeneous or heterogenous
+        :pyg:`PyG` graph.
 
         Args:
             edge_id (torch.Tensor): The global identifier for every local edge.
             edge_index (torch.Tensor): The local edge indices.
             num_nodes (int): The number of nodes in the local graph.
+            is_sorted (bool): Indicate if edge_index is sorted on col/dst_node
+                (CSC format). (default: :obj:`False`)
         """
+        graph_store = cls()
+
+        if not is_sorted:
+            edge_index, edge_id = sort_edge_index(
+                edge_index,
+                edge_id,
+                sort_by_row=False,
+            )
+
         attr = dict(
             edge_type=None,
             layout='coo',
             size=(num_nodes, num_nodes),
+            is_sorted=True,
         )
 
-        graph_store = cls()
         graph_store.put_edge_index(edge_index, **attr)
         graph_store.put_edge_id(edge_id, **attr)
+
         return graph_store
 
     @classmethod
@@ -115,7 +124,8 @@ class LocalGraphStore(GraphStore):
         edge_id_dict: Dict[EdgeType, Tensor],
         edge_index_dict: Dict[EdgeType, Tensor],
         num_nodes_dict: Dict[NodeType, int],
-    ) -> 'LocalGraphStore':
+        is_sorted: bool = False,
+    ) -> "LocalGraphStore":
         r"""Creates a local graph store from a heterogeneous :pyg:`PyG` graph.
 
         Args:
@@ -123,22 +133,29 @@ class LocalGraphStore(GraphStore):
                 for every local edge of every edge type.
             edge_index_dict (Dict[EdgeType, torch.Tensor]): The local edge
                 indices of every edge type.
-            num_nodes_dict (Dict[NodeType, int]): The number of nodes in the
-                local graph of every node type.
+            num_nodes_dict: (Dict[str, int]): The number of nodes for every
+                node type.
+            is_sorted (bool): Indicate if edge_index is sorted on col/dst_node
+                (CSC format)
         """
-        attr_dict = {}
-        for edge_type in edge_index_dict.keys():
+        graph_store = cls()
+        graph_store.meta = {'is_hetero': True}
+
+        for edge_type, edge_index in edge_index_dict.items():
             src, _, dst = edge_type
-            attr_dict[edge_type] = dict(
+            attr = dict(
                 edge_type=edge_type,
                 layout='coo',
                 size=(num_nodes_dict[src], num_nodes_dict[dst]),
+                is_sorted=True,
             )
-
-        graph_store = cls()
-        for edge_type, edge_index in edge_index_dict.items():
-            attr = attr_dict[edge_type]
             edge_id = edge_id_dict[edge_type]
+            if not is_sorted:
+                edge_index, edge_id = sort_edge_index(
+                    edge_index,
+                    edge_id,
+                    sort_by_row=False,
+                )
             graph_store.put_edge_index(edge_index, **attr)
             graph_store.put_edge_id(edge_id, **attr)
         return graph_store
@@ -152,22 +169,37 @@ class LocalGraphStore(GraphStore):
         assert osp.exists(part_dir)
 
         graph_data = torch.load(osp.join(part_dir, 'graph.pt'))
-
         graph_store = cls()
+        graph_store.is_sorted = meta['is_sorted']
 
         if not meta['is_hetero']:
-            attr = dict(edge_type=None, layout='coo', size=graph_data['size'])
-            graph_store.put_edge_index(
-                torch.stack((graph_data['row'], graph_data['col']), dim=0),
-                **attr)
-            graph_store.put_edge_id(graph_data['edge_id'], **attr)
+            edge_index = torch.stack((graph_data['row'], graph_data['col']),
+                                     dim=0)
+            edge_id = graph_data['edge_id']
+            if not graph_store.is_sorted:
+                edge_index, edge_id = sort_edge_index(edge_index, edge_id,
+                                                      sort_by_row=False)
+
+            attr = dict(edge_type=None, layout='coo', size=graph_data['size'],
+                        is_sorted=True)
+            graph_store.put_edge_index(edge_index, **attr)
+            graph_store.put_edge_id(edge_id, **attr)
 
         if meta['is_hetero']:
             for edge_type, data in graph_data.items():
-                attr = dict(edge_type=edge_type, layout='coo',
-                            size=data['size'])
-                graph_store.put_edge_index(
-                    torch.stack((data['row'], data['col']), dim=0), **attr)
-                graph_store.put_edge_id(data['edge_id'], **attr)
+                attr = dict(
+                    edge_type=edge_type,
+                    layout='coo',
+                    size=data['size'],
+                    is_sorted=True,
+                )
+                edge_index = torch.stack((data['row'], data['col']), dim=0)
+                edge_id = data['edge_id']
+
+                if not graph_store.is_sorted:
+                    edge_index, edge_id = sort_edge_index(
+                        edge_index, edge_id, sort_by_row=False)
+                graph_store.put_edge_index(edge_index, **attr)
+                graph_store.put_edge_id(edge_id, **attr)
 
         return graph_store
diff --git a/torch_geometric/distributed/partition.py b/torch_geometric/distributed/partition.py
index 73b5a9345..7c767f13f 100644
--- a/torch_geometric/distributed/partition.py
+++ b/torch_geometric/distributed/partition.py
@@ -7,8 +7,10 @@ from typing import List, Optional, Union
 import torch
 
 from torch_geometric.data import Data, HeteroData
-from torch_geometric.loader import ClusterData
-from torch_geometric.typing import EdgeType, EdgeTypeStr, NodeType
+from torch_geometric.distributed.utils import as_str
+from torch_geometric.loader.cluster import ClusterData
+from torch_geometric.typing import Dict, EdgeType, EdgeTypeStr, NodeType, Tuple
+from torch_geometric.utils import index_sort
 
 
 class Partitioner:
@@ -94,13 +96,13 @@ class Partitioner:
     def generate_partition(self):
         r"""Generates the partition."""
         os.makedirs(self.root, exist_ok=True)
-
         logging.info('Saving metadata')
         meta = {
             'num_parts': self.num_parts,
             'is_hetero': self.is_hetero,
             'node_types': self.node_types,
             'edge_types': self.edge_types,
+            'is_sorted': True,  # Based on col/destination.
         }
         with open(osp.join(self.root, 'META.json'), 'w') as f:
             json.dump(meta, f)
@@ -151,45 +153,56 @@ class Partitioner:
                 node_id = node_perm[start:end]
                 node_map[node_id] = pid
 
-                out = {}
+                graph = {}
+                efeat = {}
                 for i, edge_type in enumerate(self.edge_types):
                     src, _, dst = edge_type
                     size = (self.data[src].num_nodes, self.data[dst].num_nodes)
 
                     mask = part_data.edge_type == i
                     rows = part_data.edge_index[0, mask]
-                    cols = part_data.edge_index[1, mask]
-                    global_rows = node_id[rows]
-                    global_cols = node_perm[cols]
-                    out[edge_type] = {
-                        'edge_id': edge_id[mask],
-                        'row': global_rows,
-                        'col': global_cols,
+                    col = part_data.edge_index[1, mask]
+                    global_row = node_id[rows]
+                    global_col = node_perm[col]
+
+                    # Sort on col to avoid keeping track of permuations in
+                    # NeighborSampler when converting to CSC format:
+                    num_cols = col.size()[0]
+                    global_col, perm = index_sort(global_col,
+                                                  max_value=num_cols)
+                    global_row = global_row[perm]
+                    eid = edge_id[mask][perm]
+                    assert torch.equal(
+                        data.edge_index[:, eid],
+                        torch.stack((global_row, global_col), dim=0))
+
+                    graph[edge_type] = {
+                        'edge_id': eid,
+                        'row': global_row,
+                        'col': global_col,
                         'size': size,
                     }
-                torch.save(out, osp.join(path, 'graph.pt'))
 
-                out = {}
+                    if 'edge_attr' in part_data:
+                        edge_attr = part_data.edge_attr[mask][perm]
+                        assert torch.equal(data.edge_attr[eid, :], edge_attr)
+                        efeat[edge_type] = {
+                            'global_id': eid,
+                            'feats': dict(edge_attr=edge_attr),
+                        }
+
+                torch.save(efeat, osp.join(path, 'edge_feats.pt'))
+                torch.save(graph, osp.join(path, 'graph.pt'))
+
+                nfeat = {}
                 for i, node_type in enumerate(self.node_types):
                     mask = part_data.node_type == i
                     x = part_data.x[mask] if 'x' in part_data else None
-                    out[node_type] = {
+                    nfeat[node_type] = {
                         'global_id': node_id[mask],
                         'feats': dict(x=x),
                     }
-                torch.save(out, osp.join(path, 'node_feats.pt'))
-
-                out = {}
-                for i, edge_type in enumerate(self.edge_types):
-                    mask = part_data.edge_type == i
-                    edge_attr = None
-                    if 'edge_attr' in part_data:
-                        edge_attr = part_data.edge_attr[mask]
-                    out[edge_type] = {
-                        'global_id': edge_id[mask],
-                        'feats': dict(edge_attr=edge_attr),
-                    }
-                torch.save(out, osp.join(path, 'edge_feats.pt'))
+                torch.save(nfeat, osp.join(path, 'node_feats.pt'))
 
             logging.info('Saving partition mapping info')
 
@@ -207,7 +220,6 @@ class Partitioner:
                            osp.join(path, f'{EdgeTypeStr(edge_type)}.pt'))
 
         else:  # `if not self.is_hetero:`
-
             edge_start = 0
             for pid in range(self.num_parts):
                 logging.info(f'Saving graph partition {pid}')
@@ -222,18 +234,34 @@ class Partitioner:
                 edge_map[edge_id] = pid
                 edge_start += num_edges
 
-                node_id = node_perm[start:end]
-                node_map[node_id] = pid
+                node_id = node_perm[start:end]  # global node_ids
+                node_map[node_id] = pid  # 0 or 1
+
                 rows = part_data.edge_index[0]
-                cols = part_data.edge_index[1]
-                global_rows = node_id[rows]
-                global_cols = node_perm[cols]
+                col = part_data.edge_index[1]
+                num_cols = col.size()[0]
+
+                global_row = node_id[rows]  # part_ids -> global
+                global_col = node_perm[col]
+
+                # Sort on col to avoid keeping track of permuations in
+                # NeighborSampler when converting to CSC format:
+                global_col, perm = index_sort(global_col, max_value=num_cols)
+                global_row = global_row[perm]
+                edge_id = edge_id[perm]
+
+                assert torch.equal(self.data.edge_index[:, edge_id],
+                                   torch.stack((global_row, global_col)))
+                if 'edge_attr' in part_data:
+                    edge_attr = part_data.edge_attr[perm]
+                    assert torch.equal(self.data.edge_attr[edge_id, :],
+                                       edge_attr)
 
                 torch.save(
                     {
                         'edge_id': edge_id,
-                        'row': global_rows,
-                        'col': global_cols,
+                        'row': global_row,
+                        'col': global_col,
                         'size': (data.num_nodes, data.num_nodes),
                     }, osp.join(path, 'graph.pt'))
 
@@ -242,13 +270,49 @@ class Partitioner:
                         'global_id': node_id,
                         'feats': dict(x=part_data.x),
                     }, osp.join(path, 'node_feats.pt'))
-
-                torch.save(
-                    {
-                        'global_id': edge_id,
-                        'feats': dict(edge_attr=part_data.edge_attr),
-                    }, osp.join(path, 'edge_feats.pt'))
+                if 'edge_attr' in part_data:
+                    torch.save(
+                        {
+                            'global_id': edge_id,
+                            'feats': dict(edge_attr=part_data.edge_attr[perm]),
+                        }, osp.join(path, 'edge_feats.pt'))
 
             logging.info('Saving partition mapping info')
             torch.save(node_map, osp.join(self.root, 'node_map.pt'))
             torch.save(edge_map, osp.join(self.root, 'edge_map.pt'))
+
+
+def load_partition_info(
+    root_dir: str,
+    partition_idx: int,
+) -> Tuple[Dict, int, int, torch.Tensor, torch.Tensor]:
+
+    # load the partition with PyG format (graphstore/featurestore)
+    with open(os.path.join(root_dir, 'META.json'), 'rb') as infile:
+        meta = json.load(infile)
+    num_partitions = meta['num_parts']
+    assert partition_idx >= 0
+    assert partition_idx < num_partitions
+    partition_dir = os.path.join(root_dir, f'part_{partition_idx}')
+    assert os.path.exists(partition_dir)
+
+    if meta['is_hetero'] is False:
+        node_pb = torch.load(os.path.join(root_dir, 'node_map.pt'))
+        edge_pb = torch.load(os.path.join(root_dir, 'edge_map.pt'))
+
+        return (meta, num_partitions, partition_idx, node_pb, edge_pb)
+    else:
+        node_pb_dict = {}
+        node_pb_dir = os.path.join(root_dir, 'node_map')
+        for ntype in meta['node_types']:
+            node_pb_dict[ntype] = torch.load(
+                os.path.join(node_pb_dir, f'{as_str(ntype)}.pt'))
+
+        edge_pb_dict = {}
+        edge_pb_dir = os.path.join(root_dir, 'edge_map')
+        for etype in meta['edge_types']:
+            edge_pb_dict[tuple(etype)] = torch.load(
+                os.path.join(edge_pb_dir, f'{as_str(etype)}.pt'))
+
+        return (meta, num_partitions, partition_idx, node_pb_dict,
+                edge_pb_dict)
diff --git a/torch_geometric/distributed/utils.py b/torch_geometric/distributed/utils.py
index d2adf9339..babb65aaf 100644
--- a/torch_geometric/distributed/utils.py
+++ b/torch_geometric/distributed/utils.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Tuple
+from typing import Dict, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -8,7 +8,7 @@ from torch import Tensor
 from torch_geometric.data import HeteroData
 from torch_geometric.distributed import LocalFeatureStore, LocalGraphStore
 from torch_geometric.sampler import SamplerOutput
-from torch_geometric.typing import Dict, NodeType, Optional
+from torch_geometric.typing import EdgeType, NodeType
 
 
 @dataclass
@@ -80,6 +80,7 @@ def filter_dist_store(
     edge_dict: Dict[str, Optional[Tensor]],
     custom_cls: Optional[HeteroData] = None,
     meta: Optional[Dict[str, Tensor]] = None,
+    input_type: str = None,
 ) -> HeteroData:
     r"""Constructs a :class:`HeteroData` object from a feature store that only
     holds nodes in `node` end edges in `edge` for each node and edge type,
@@ -117,7 +118,27 @@ def filter_dist_store(
             if efeats[attr.edge_type] is not None:
                 data[attr.edge_type].edge_attr = efeats[attr.edge_type]
 
-    for label in nlabels:
-        data[label].y = nlabels[label]
+    if nlabels is not None:
+        data[input_type].y = nlabels
 
     return data
+
+
+def as_str(inputs: Union[NodeType, EdgeType]) -> str:
+    if isinstance(inputs, NodeType):
+        return inputs
+    elif isinstance(inputs, (list, tuple)) and len(inputs) == 3:
+        return '__'.join(inputs)
+    return ''
+
+
+def reverse_edge_type(etype: EdgeType) -> EdgeType:
+    src, rel, dst = etype
+    if src != dst:
+        if rel.split('_', 1)[0] == 'rev':
+            # undirected edge with `rev_` prefix.
+            rel = rel.split('_', 1)[1]
+        else:
+            rel = 'rev_' + rel
+
+    return dst, rel, src


Test patch:
diff --git a/test/distributed/test_local_graph_store.py b/test/distributed/test_local_graph_store.py
index ce83082af..fc36536c3 100644
--- a/test/distributed/test_local_graph_store.py
+++ b/test/distributed/test_local_graph_store.py
@@ -10,17 +10,25 @@ def test_local_graph_store():
     edge_index = get_random_edge_index(100, 100, 300)
     edge_id = torch.tensor([1, 2, 3, 5, 8, 4])
 
-    graph_store.put_edge_index(edge_index, edge_type=None, layout='coo',
-                               size=(100, 100))
+    graph_store.put_edge_index(
+        edge_index,
+        edge_type=None,
+        layout='coo',
+        size=(100, 100),
+    )
 
-    graph_store.put_edge_id(edge_id, edge_type=None, layout='coo',
-                            size=(100, 100))
+    graph_store.put_edge_id(
+        edge_id,
+        edge_type=None,
+        layout='coo',
+        size=(100, 100),
+    )
 
     assert len(graph_store.get_all_edge_attrs()) == 1
     edge_attr = graph_store.get_all_edge_attrs()[0]
     assert torch.equal(graph_store.get_edge_index(edge_attr), edge_index)
     assert torch.equal(graph_store.get_edge_id(edge_attr), edge_id)
-
+    assert not graph_store.is_sorted
     graph_store.remove_edge_index(edge_attr)
     graph_store.remove_edge_id(edge_attr)
     assert len(graph_store.get_all_edge_attrs()) == 0
@@ -29,14 +37,20 @@ def test_local_graph_store():
 def test_homogeneous_graph_store():
     edge_id = torch.randperm(300)
     edge_index = get_random_edge_index(100, 100, 300)
+    edge_index[1] = torch.sort(edge_index[1])[0]
 
-    graph_store = LocalGraphStore.from_data(edge_id, edge_index, num_nodes=100)
+    graph_store = LocalGraphStore.from_data(
+        edge_id,
+        edge_index,
+        num_nodes=100,
+        is_sorted=True,
+    )
 
     assert len(graph_store.get_all_edge_attrs()) == 1
     edge_attr = graph_store.get_all_edge_attrs()[0]
     assert edge_attr.edge_type is None
     assert edge_attr.layout.value == 'coo'
-    assert not edge_attr.is_sorted
+    assert edge_attr.is_sorted
     assert edge_attr.size == (100, 100)
 
     assert torch.equal(
@@ -52,16 +66,22 @@ def test_homogeneous_graph_store():
 def test_heterogeneous_graph_store():
     edge_type = ('paper', 'to', 'paper')
     edge_id_dict = {edge_type: torch.randperm(300)}
-    edge_index_dict = {edge_type: get_random_edge_index(100, 100, 300)}
+    edge_index = get_random_edge_index(100, 100, 300)
+    edge_index[1] = torch.sort(edge_index[1])[0]
+    edge_index_dict = {edge_type: edge_index}
 
     graph_store = LocalGraphStore.from_hetero_data(
-        edge_id_dict, edge_index_dict, num_nodes_dict={'paper': 100})
+        edge_id_dict,
+        edge_index_dict,
+        num_nodes_dict={'paper': 100},
+        is_sorted=True,
+    )
 
     assert len(graph_store.get_all_edge_attrs()) == 1
     edge_attr = graph_store.get_all_edge_attrs()[0]
     assert edge_attr.edge_type == edge_type
     assert edge_attr.layout.value == 'coo'
-    assert not edge_attr.is_sorted
+    assert edge_attr.is_sorted
     assert edge_attr.size == (100, 100)
 
     assert torch.equal(
@@ -72,3 +92,45 @@ def test_heterogeneous_graph_store():
         graph_store.get_edge_index(edge_type, layout='coo'),
         edge_index_dict[edge_type],
     )
+
+
+def test_sorted_graph_store():
+    edge_index_sorted = torch.tensor([[1, 7, 5, 6, 1], [0, 0, 1, 1, 2]])
+    edge_id_sorted = torch.tensor([0, 1, 2, 3, 4])
+
+    edge_index = torch.tensor([[1, 5, 7, 1, 6], [0, 1, 0, 2, 1]])
+    edge_id = torch.tensor([0, 2, 1, 4, 3])
+
+    graph_store = LocalGraphStore.from_data(
+        edge_id,
+        edge_index,
+        num_nodes=8,
+        is_sorted=False,
+    )
+    assert torch.equal(
+        graph_store.get_edge_index(edge_type=None, layout='coo'),
+        edge_index_sorted,
+    )
+    assert torch.equal(
+        graph_store.get_edge_id(edge_type=None, layout='coo'),
+        edge_id_sorted,
+    )
+
+    edge_type = ('paper', 'to', 'paper')
+    edge_index_dict = {edge_type: edge_index}
+    edge_id_dict = {edge_type: edge_id}
+
+    graph_store = LocalGraphStore.from_hetero_data(
+        edge_id_dict,
+        edge_index_dict,
+        num_nodes_dict={'paper': 8},
+        is_sorted=False,
+    )
+    assert torch.equal(
+        graph_store.get_edge_index(edge_type, layout='coo'),
+        edge_index_sorted,
+    )
+    assert torch.equal(
+        graph_store.get_edge_id(edge_type, layout='coo'),
+        edge_id_sorted,
+    )