Repository: dmlc/dgl

Issue:
Number of nodes wrong when batching blocks
## 🐛 Bug

See the following reproduction code

## To Reproduce

```python
plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
wishes_nx = nx.DiGraph()
wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1)
wishes_nx.add_edge('u0', 'g1', id=0)
wishes_nx.add_edge('u2', 'g0', id=1)

follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game')
wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game')
develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game')
g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
g2 = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])

block1 = dgl.to_block(g, {'user': [1, 2], 'game': [0, 1], 'developer': []})
block2 = dgl.to_block(g2, {'user': [1, 2], 'game': [0, 1], 'developer': []})
batched_block = dgl.batch_hetero([block1, block2])
assert block1.number_of_src_nodes('developer') == 2             # passes
assert block2.number_of_src_nodes('developer') == 2             # passes
assert batched_block.number_of_src_nodes('developer') == 4      # fails, = 0
```

## Expected behavior

The last test should return pass.

## Environment

 - DGL Version (e.g., 1.0): master
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3):
 - OS (e.g., Linux):
 - How you installed DGL (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:



Discussion:
It seems batching of block is not well tested
I also came across a batching problem, directly with heterographs:
```
g0 = dgl.heterograph({
     ('user', 'plays', 'game'): [(0, 0), (1, 0), (1, 1), (2, 1)],
     }, num_nodes_dict={"foo":1, "user":3, "game":2})
g1 = dgl.heterograph({
     ('user', 'plays', 'game'): [(0, 0), (1, 0), (1, 1), (2, 1)],
     }, num_nodes_dict={"foo": 2, "user":3, "game":2})
gall = dgl.batch_hetero([g0, g1])
print(g0)
print(g1)
print(gall)
```
prints:
```
Graph(num_nodes={'foo': 1, 'game': 2, 'user': 3},
      num_edges={('user', 'plays', 'game'): 4},
      metagraph=[('user', 'game')])
Graph(num_nodes={'foo': 2, 'game': 2, 'user': 3},
      num_edges={('user', 'plays', 'game'): 4},
      metagraph=[('user', 'game')])
Graph(num_nodes={'foo': 0, 'game': 4, 'user': 6},
      num_edges={('user', 'plays', 'game'): 8},
      metagraph=[('user', 'game')])
```
So it seems that nodes without edges get dropped when batching heterographs

Resolving patch:
diff --git a/python/dgl/batch.py b/python/dgl/batch.py
--- a/python/dgl/batch.py
+++ b/python/dgl/batch.py
@@ -1,5 +1,6 @@
 """Utilities for batching/unbatching graphs."""
 from collections.abc import Mapping
+from collections import defaultdict
 
 from . import backend as F
 from .base import ALL, is_all, DGLError, dgl_warning
@@ -167,24 +168,25 @@ def batch(graphs, ndata=ALL, edata=ALL, *, node_attrs=None, edge_attrs=None):
     utils.check_all_same_device(graphs, 'graphs')
     utils.check_all_same_idtype(graphs, 'graphs')
     relations = graphs[0].canonical_etypes
+    ntypes = graphs[0].ntypes
     idtype = graphs[0].idtype
     device = graphs[0].device
 
     # Batch graph structure for each relation graph
-    edge_dict = {}
-    num_nodes_dict = {}
-    for rel in relations:
-        srctype, etype, dsttype = rel
-        srcnid_off = dstnid_off = 0
-        src, dst = [], []
-        for g in graphs:
+    edge_dict = defaultdict(list)
+    num_nodes_dict = defaultdict(int)
+    for g in graphs:
+        for rel in relations:
+            srctype, etype, dsttype = rel
             u, v = g.edges(order='eid', etype=rel)
-            src.append(u + srcnid_off)
-            dst.append(v + dstnid_off)
-            srcnid_off += g.number_of_nodes(srctype)
-            dstnid_off += g.number_of_nodes(dsttype)
+            src = u + num_nodes_dict[srctype]
+            dst = v + num_nodes_dict[dsttype]
+            edge_dict[rel].append((src, dst))
+        for ntype in ntypes:
+            num_nodes_dict[ntype] += g.number_of_nodes(ntype)
+    for rel in relations:
+        src, dst = zip(*edge_dict[rel])
         edge_dict[rel] = (F.cat(src, 0), F.cat(dst, 0))
-        num_nodes_dict.update({srctype : srcnid_off, dsttype : dstnid_off})
     retg = convert.heterograph(edge_dict, num_nodes_dict, idtype=idtype, device=device)
 
     # Compute batch num nodes


Test patch:
diff --git a/test.py b/test.py
deleted file mode 100644
diff --git a/tests/compute/test_batched_heterograph.py b/tests/compute/test_batched_heterograph.py
--- a/tests/compute/test_batched_heterograph.py
+++ b/tests/compute/test_batched_heterograph.py
@@ -1,10 +1,11 @@
 import dgl
 import backend as F
 import unittest
+import pytest
 
 from dgl.base import ALL
 from utils import parametrize_dtype
-from test_utils import check_graph_equal
+from test_utils import check_graph_equal, get_cases
 
 def check_equivalence_between_heterographs(g1, g2, node_attrs=None, edge_attrs=None):
     assert g1.ntypes == g2.ntypes
@@ -40,19 +41,13 @@ def check_equivalence_between_heterographs(g1, g2, node_attrs=None, edge_attrs=N
             for feat_name in edge_attrs[ety]:
                 assert F.allclose(g1.edges[ety].data[feat_name], g2.edges[ety].data[feat_name])
 
+@pytest.mark.parametrize('gs', get_cases(['two_hetero_batch']))
 @parametrize_dtype
-def test_topology(idtype):
+def test_topology(gs, idtype):
     """Test batching two DGLHeteroGraphs where some nodes are isolated in some relations"""
-    g1 = dgl.heterograph({
-        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
-        ('user', 'follows', 'developer'): ([0, 1], [1, 2]),
-        ('user', 'plays', 'game'): ([0, 1, 2, 3], [0, 0, 1, 1])
-    }, idtype=idtype, device=F.ctx())
-    g2 = dgl.heterograph({
-        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
-        ('user', 'follows', 'developer'): ([0, 1], [1, 2]),
-        ('user', 'plays', 'game'): ([0, 1, 2], [0, 0, 1])
-    }, idtype=idtype, device=F.ctx())
+    g1, g2 = gs
+    g1 = g1.astype(idtype).to(F.ctx())
+    g2 = g2.astype(idtype).to(F.ctx())
     bg = dgl.batch([g1, g2])
 
     assert bg.idtype == idtype
diff --git a/tests/compute/test_pickle.py b/tests/compute/test_pickle.py
--- a/tests/compute/test_pickle.py
+++ b/tests/compute/test_pickle.py
@@ -132,7 +132,7 @@ def _global_message_func(nodes):
 
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
 @parametrize_dtype
-@pytest.mark.parametrize('g', get_cases(exclude=['dglgraph']))
+@pytest.mark.parametrize('g', get_cases(exclude=['dglgraph', 'two_hetero_batch']))
 def test_pickling_graph(g, idtype):
     g = g.astype(idtype)
     new_g = _reconstruct_pickle(g)
diff --git a/tests/test_utils/graph_cases.py b/tests/test_utils/graph_cases.py
--- a/tests/test_utils/graph_cases.py
+++ b/tests/test_utils/graph_cases.py
@@ -109,3 +109,31 @@ def random_bipartite(size_src, size_dst):
 def random_block(size):
     g = dgl.from_networkx(nx.erdos_renyi_graph(size, 0.1))
     return dgl.to_block(g, np.unique(F.zerocopy_to_numpy(g.edges()[1])))
+
+@register_case(['two_hetero_batch'])
+def two_hetero_batch():
+    g1 = dgl.heterograph({
+        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
+        ('user', 'follows', 'developer'): ([0, 1], [1, 2]),
+        ('user', 'plays', 'game'): ([0, 1, 2, 3], [0, 0, 1, 1])
+    })
+    g2 = dgl.heterograph({
+        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
+        ('user', 'follows', 'developer'): ([0, 1], [1, 2]),
+        ('user', 'plays', 'game'): ([0, 1, 2], [0, 0, 1])
+    })
+    return [g1, g2]
+
+@register_case(['two_hetero_batch'])
+def two_hetero_batch_with_isolated_ntypes():
+    g1 = dgl.heterograph({
+        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
+        ('user', 'follows', 'developer'): ([0, 1], [1, 2]),
+        ('user', 'plays', 'game'): ([0, 1, 2, 3], [0, 0, 1, 1])
+        }, num_nodes_dict={'user': 4, 'game': 2, 'developer': 3, 'platform': 2})
+    g2 = dgl.heterograph({
+        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
+        ('user', 'follows', 'developer'): ([0, 1], [1, 2]),
+        ('user', 'plays', 'game'): ([0, 1, 2], [0, 0, 1])
+        }, num_nodes_dict={'user': 3, 'game': 2, 'developer': 3, 'platform': 3})
+    return [g1, g2]