Repository: pyg-team/pytorch_geometric
Pull request #8669 (2023-12-23 11:45:54)

Support for `EdgeIndex` in `DataLoader`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8669?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`0a89cc1`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/0a89cc105ad55aea1c466a404a2fffbcdadee0de?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.45% compared to head [(`4495850`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8669?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.45%.

> :exclamation: Current head 4495850 differs from pull request most recent head 8b9d8c0. Consider uploading reports for the commit 8b9d8c0 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/data/storage.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvc3RvcmFnZS5weQ==) | 66.66% | [2 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #8669      +/-   ##
==========================================
- Coverage   89.45%   89.45%   -0.01%     
==========================================
  Files         481      481              
  Lines       30834    30847      +13     
==========================================
+ Hits        27583    27594      +11     
- Misses       3251     3253       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8669?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/torch_geometric/data/collate.py b/torch_geometric/data/collate.py
index baba7c473..474d6e1cf 100644
--- a/torch_geometric/data/collate.py
+++ b/torch_geometric/data/collate.py
@@ -6,8 +6,10 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
+from torch_geometric import EdgeIndex
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage, NodeStorage
+from torch_geometric.edge_index import SortOrder
 from torch_geometric.typing import (
     SparseTensor,
     TensorFrame,
@@ -176,6 +178,14 @@ def _collate(
 
         value = torch.cat(values, dim=cat_dim or 0, out=out)
 
+        if increment and isinstance(value, EdgeIndex) and values[0].is_sorted:
+            # Check whether the whole `EdgeIndex` is sorted by row:
+            if values[0].is_sorted_by_row and (value[0].diff() >= 0).all():
+                value._sort_order = SortOrder.ROW
+            # Check whether the whole `EdgeIndex` is sorted by column:
+            elif values[0].is_sorted_by_col and (value[1].diff() >= 0).all():
+                value._sort_order = SortOrder.COL
+
         return value, slices, incs
 
     elif isinstance(elem, TensorFrame):
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index 1096f5d33..1a815a3a7 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -24,6 +24,7 @@ import torch
 from torch import Tensor
 from typing_extensions import Self
 
+from torch_geometric import EdgeIndex
 from torch_geometric.data.view import ItemsView, KeysView, ValuesView
 from torch_geometric.typing import (
     EdgeType,
@@ -436,6 +437,11 @@ class NodeStorage(BaseStorage):
                 return value.shape[cat_dim]
             if isinstance(value, TensorFrame) and 'node' in key:
                 return value.num_rows
+        if 'edge_index' in self and isinstance(self.edge_index, EdgeIndex):
+            if self.edge_index.sparse_size(0) is not None:
+                return self.edge_index.sparse_size(0)
+            if self.edge_index.sparse_size(1) is not None:
+                return self.edge_index.sparse_size(1)
         if 'adj' in self and isinstance(self.adj, SparseTensor):
             return self.adj.size(0)
         if 'adj_t' in self and isinstance(self.adj_t, SparseTensor):


Test patch:
diff --git a/test/loader/test_dataloader.py b/test/loader/test_dataloader.py
index 2ccd8251e..6017fcaf1 100644
--- a/test/loader/test_dataloader.py
+++ b/test/loader/test_dataloader.py
@@ -5,6 +5,7 @@ from collections import namedtuple
 import pytest
 import torch
 
+from torch_geometric import EdgeIndex
 from torch_geometric.data import Data, HeteroData, OnDiskDataset
 from torch_geometric.loader import DataLoader
 from torch_geometric.testing import (
@@ -186,6 +187,37 @@ def test_heterogeneous_dataloader(num_workers):
             assert id(batch) == id(store._parent())
 
 
+@pytest.mark.parametrize('num_workers', num_workers_list)
+@pytest.mark.parametrize('sort_order', [None, 'row', 'col'])
+def test_edge_index_dataloader(num_workers, sort_order):
+    if sort_order == 'col':
+        edge_index = [[1, 0, 2, 1], [0, 1, 1, 2]]
+    else:
+        edge_index = [[0, 1, 1, 2], [1, 0, 2, 1]]
+
+    edge_index = EdgeIndex(
+        edge_index,
+        sparse_size=(3, 3),
+        sort_order=sort_order,
+        is_undirected=True,
+    )
+    data = Data(edge_index=edge_index)
+    assert data.num_nodes == 3
+
+    loader = DataLoader(
+        [data, data, data, data],
+        batch_size=2,
+        num_workers=num_workers,
+    )
+    assert len(loader) == 2
+
+    for batch in loader:
+        assert isinstance(batch.edge_index, EdgeIndex)
+        assert batch.edge_index.sparse_size() == (6, 6)
+        assert batch.edge_index.sort_order == sort_order
+        assert batch.edge_index.is_undirected
+
+
 @withPackage('torch_frame')
 def test_dataloader_tensor_frame():
     tf = get_random_tensor_frame(num_rows=10)