Repository: dmlc/dgl

Issue:
When an ntype or etype is empty in every input graph, dgl.batch() erases their data frames in the output
## 🐛 Bug

If there are no nodes or edges of a particular ntype or etype in any of the input graphs, dgl.batch() erases their data frames in the output batch.

## To Reproduce

Take graphs g1 and g2 with a node feature x for ntype t. Batch them with `b = dgl.batch([g1, g2])`. Later computations access `b.nodes["t"].data["x"]`, expecting it to exist, even if empty. However, sometimes each graph in the input happens to have no nodes of ntype t. When this happens, `b.nodes["t"].data["x"]` will not exist and it will raise a KeyError. The same happens with edges and edge features.  

## Expected behavior

Even if there are no nodes or edges of that ntype or etype containing that feature in the input graphs, the output batch should have the same features as the input graphs.

## Environment

 - DGL Version (e.g., 1.0): 0.5.3
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.6.0
 - How you installed DGL (`conda`, `pip`, source): conda
 - Python version: 3.7


Discussion:

Resolving patch:
diff --git a/python/dgl/batch.py b/python/dgl/batch.py
--- a/python/dgl/batch.py
+++ b/python/dgl/batch.py
@@ -11,7 +11,8 @@
 
 __all__ = ['batch', 'unbatch', 'batch_hetero', 'unbatch_hetero']
 
-def batch(graphs, ndata=ALL, edata=ALL, *, node_attrs=None, edge_attrs=None):
+def batch(graphs, ndata=ALL, edata=ALL, *,
+          node_attrs=None, edge_attrs=None):
     r"""Batch a collection of :class:`DGLGraph` s into one graph for more efficient
     graph computation.
 
@@ -191,9 +192,10 @@ def batch(graphs, ndata=ALL, edata=ALL, *, node_attrs=None, edge_attrs=None):
     # Batch node feature
     if ndata is not None:
         for ntype_id, ntype in zip(ntype_ids, ntypes):
+            all_empty = all(g._graph.number_of_nodes(ntype_id) == 0 for g in graphs)
             frames = [
                 g._node_frames[ntype_id] for g in graphs
-                if g._graph.number_of_nodes(ntype_id) > 0]
+                if g._graph.number_of_nodes(ntype_id) > 0 or all_empty]
             # TODO: do we require graphs with no nodes/edges to have the same schema?  Currently
             # we allow empty graphs to have no features during batching.
             ret_feat = _batch_feat_dicts(frames, ndata, 'nodes["{}"].data'.format(ntype))
@@ -202,9 +204,10 @@ def batch(graphs, ndata=ALL, edata=ALL, *, node_attrs=None, edge_attrs=None):
     # Batch edge feature
     if edata is not None:
         for etype_id, etype in zip(relation_ids, relations):
+            all_empty = all(g._graph.number_of_edges(etype_id) == 0 for g in graphs)
             frames = [
                 g._edge_frames[etype_id] for g in graphs
-                if g._graph.number_of_edges(etype_id) > 0]
+                if g._graph.number_of_edges(etype_id) > 0 or all_empty]
             # TODO: do we require graphs with no nodes/edges to have the same schema?  Currently
             # we allow empty graphs to have no features during batching.
             ret_feat = _batch_feat_dicts(frames, edata, 'edges[{}].data'.format(etype))


Test patch:
diff --git a/tests/compute/test_batched_graph.py b/tests/compute/test_batched_graph.py
--- a/tests/compute/test_batched_graph.py
+++ b/tests/compute/test_batched_graph.py
@@ -207,6 +207,18 @@ def test_batch_no_edge(idtype):
     g3.add_nodes(1)  # no edges
     g = dgl.batch([g1, g3, g2]) # should not throw an error
 
+@parametrize_dtype
+def test_batch_keeps_empty_data(idtype):
+    g1 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g1.ndata["nh"] = F.tensor([])
+    g1.edata["eh"] = F.tensor([]) 
+    g2 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g2.ndata["nh"] = F.tensor([])
+    g2.edata["eh"] = F.tensor([]) 
+    g = dgl.batch([g1, g2])
+    assert "nh" in g.ndata
+    assert "eh" in g.edata    
+
 def _get_subgraph_batch_info(keys, induced_indices_arr, batch_num_objs):
     """Internal function to compute batch information for subgraphs.
     Parameters
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
@@ -321,6 +321,18 @@ def test_unbatch2(idtype):
     check_graph_equal(g2, gg2)
     check_graph_equal(g3, gg3)
 
+@parametrize_dtype
+def test_batch_keeps_empty_data(idtype):
+    g1 = dgl.heterograph({("a", "to", "a"): ([], [])}).astype(idtype).to(F.ctx())
+    g1.nodes["a"].data["nh"] = F.tensor([])
+    g1.edges[("a", "to", "a")].data["eh"] = F.tensor([]) 
+    g2 = dgl.heterograph({("a", "to", "a"): ([], [])}).astype(idtype).to(F.ctx())
+    g2.nodes["a"].data["nh"] = F.tensor([])
+    g2.edges[("a", "to", "a")].data["eh"] = F.tensor([]) 
+    g = dgl.batch([g1, g2])
+    assert "nh" in g.nodes["a"].data
+    assert "eh" in g.edges[("a", "to", "a")].data
+
 if __name__ == '__main__':
     #test_topology('int32')
     #test_batching_batched('int32')