GraphStore: Data, HeteroData respect is_sorted

CHANGELOG.md

@@ -12,7 +12,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `time_attr` argument to `LinkNeighborLoader` ([#4877](https://github.com/pyg-team/pytorch_geometric/pull/4877), [#4908](https://github.com/pyg-team/pytorch_geometric/pull/4908))
 - Added a `filter_per_worker` argument to data loaders to allow filtering of data within sub-processes ([#4873](https://github.com/pyg-team/pytorch_geometric/pull/4873))
 - Added a `NeighborLoader` benchmark script ([#4815](https://github.com/pyg-team/pytorch_geometric/pull/4815))
-- Added support for `FeatureStore` and `GraphStore` in `NeighborLoader` ([#4817](https://github.com/pyg-team/pytorch_geometric/pull/4817), [#4851](https://github.com/pyg-team/pytorch_geometric/pull/4851), [#4854](https://github.com/pyg-team/pytorch_geometric/pull/4854), [#4856](https://github.com/pyg-team/pytorch_geometric/pull/4856), [#4857](https://github.com/pyg-team/pytorch_geometric/pull/4857), [#4882](https://github.com/pyg-team/pytorch_geometric/pull/4882), [#4883](https://github.com/pyg-team/pytorch_geometric/pull/4883))
+- Added support for `FeatureStore` and `GraphStore` in `NeighborLoader` ([#4817](https://github.com/pyg-team/pytorch_geometric/pull/4817), [#4851](https://github.com/pyg-team/pytorch_geometric/pull/4851), [#4854](https://github.com/pyg-team/pytorch_geometric/pull/4854), [#4856](https://github.com/pyg-team/pytorch_geometric/pull/4856), [#4857](https://github.com/pyg-team/pytorch_geometric/pull/4857), [#4882](https://github.com/pyg-team/pytorch_geometric/pull/4882), [#4883](https://github.com/pyg-team/pytorch_geometric/pull/4883), [#4992](https://github.com/pyg-team/pytorch_geometric/pull/4922))
 - Added a `normalize` parameter to `dense_diff_pool` ([#4847](https://github.com/pyg-team/pytorch_geometric/pull/4847))
 - Added `size=None` explanation to jittable `MessagePassing` modules in the documentation ([#4850](https://github.com/pyg-team/pytorch_geometric/pull/4850))
 - Added documentation to the `DataLoaderIterator` class ([#4838](https://github.com/pyg-team/pytorch_geometric/pull/4838))

torch_geometric/data/graph_store.py

@@ -132,8 +132,12 @@ def get_edge_index(self, *args, **kwargs) -> EdgeTensorType:
     # Layout Conversion #######################################################
 
     # TODO support `replace` to replace the existing edge index.
-    def _to_layout(self, layout: EdgeLayout,
-                   store: bool = False) -> ConversionOutputType:
+    def _to_layout(
+        self,
+        layout: EdgeLayout,
+        store: bool = False,
+        is_sorted: bool = False,
+    ) -> ConversionOutputType:
         # Obtain all edge attributes, grouped by type:
         edge_attrs = self.get_all_edge_attrs()
         edge_type_to_attrs: Dict[Any, List[EdgeAttr]] = defaultdict(list)
@@ -200,6 +204,8 @@ def _to_layout(self, layout: EdgeLayout,
                     row, col, perm = to_csc(adj, from_attr_copy, device='cpu')
 
                 else:
+                    # Respect is_sorted override (sorted by col):
+                    from_attr.is_sorted = from_attr.is_sorted or is_sorted
                     adj = edge_tensor_type_to_adj_type(from_attr, from_tuple)
 
                     # Actually colptr, row, perm
@@ -242,10 +248,11 @@ def csr(self, store: bool = False) -> ConversionOutputType:
         optionally storing the converted edge indices in the graph store."""
         return self._to_layout(EdgeLayout.CSR, store)
 
-    def csc(self, store: bool = False) -> ConversionOutputType:
+    def csc(self, store: bool = False,
+            is_sorted: bool = False) -> ConversionOutputType:
         r"""Converts the edge indices in the graph store to CSC format,
         optionally storing the converted edge indices in the graph store."""
-        return self._to_layout(EdgeLayout.CSC, store)
+        return self._to_layout(EdgeLayout.CSC, store, is_sorted)
 
     # Additional methods ######################################################
 

torch_geometric/loader/neighbor_loader.py

@@ -122,7 +122,8 @@ def __init__(
             self.input_type = input_type
 
             # Obtain CSC representations for in-memory sampling:
-            row_dict, colptr_dict, perm_dict = graph_store.csc()
+            row_dict, colptr_dict, perm_dict = graph_store.csc(
+                is_sorted=is_sorted)
             self.row_dict = {
                 edge_type_to_str(k): v
                 for k, v in row_dict.items()