Repository: dmlc/dgl

Issue:
bug in DGLHeteroGraph edge_type_subgraph
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
I found a bug when using `edge_type_subgraph`. It expects node types to be sequential and starting from zero, when asking for a node type with high idx it breaks.

The test doesn't catch that as the indexes of the node types in the test are 0,1
## To Reproduce

Steps to reproduce the behavior:

Using the example from the tests:
```python
import dgl
from scipy import sparse as ssp
import networkx as nx

def create_test_heterograph():
    # test heterograph from the docstring, plus a user -- wishes -- game relation
    # 3 users, 2 games, 2 developers
    # metagraph:
    #    ('user', 'follows', 'user'),
    #    ('user', 'plays', 'game'),
    #    ('user', 'wishes', 'game'),
    #    ('developer', 'develops', 'game')])

    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])
    return g

g = create_test_heterograph()
g.edge_type_subgraph(['develops'])
```

```
  File "/home/doron/.PyCharm2019.2/config/scratches/scratch.py", line 29, in <module>
    g.edge_type_subgraph(['develops'])
  File "/home/doron/miniconda3/envs/DeepLearning_Projects/lib/python3.7/site-packages/dgl/heterograph.py", line 1756, in edge_type_subgraph
    mapped_meta_src = [induced_ntype_ids[v] for v in meta_src]
  File "/home/doron/miniconda3/envs/DeepLearning_Projects/lib/python3.7/site-packages/dgl/heterograph.py", line 1756, in <listcomp>
    mapped_meta_src = [induced_ntype_ids[v] for v in meta_src]
IndexError: list index out of range
```

Here for example `v=2` while `len(induced_ntype_ids)==1`
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
subgraph returned

## Environment

 - DGL Version - dgl-cuda10.0 v0.4.1 (from conda)
 - OS - Ubuntu 18.04
 - Python version: 3.7.5

## Additional context

### Suggestion for fix:

This can be easily fixed with changing:
```python
induced_ntype_ids = list(set(meta_src) | set(meta_dst))
```

to:
```python
induced_ntype_ids = {x:i for i,x in enumerate(set(meta_src) | set(meta_dst))}
```
<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/python/dgl/heterograph.py b/python/dgl/heterograph.py
--- a/python/dgl/heterograph.py
+++ b/python/dgl/heterograph.py
@@ -1752,12 +1752,12 @@ def edge_type_subgraph(self, etypes):
         rel_graphs = [self._graph.get_relation_graph(i) for i in etype_ids]
         meta_src = meta_src.tonumpy()
         meta_dst = meta_dst.tonumpy()
-        induced_ntype_ids = list(set(meta_src) | set(meta_dst))
-        mapped_meta_src = [induced_ntype_ids[v] for v in meta_src]
-        mapped_meta_dst = [induced_ntype_ids[v] for v in meta_dst]
-        node_frames = [self._node_frames[i] for i in induced_ntype_ids]
+        ntypes_invmap = {n: i for i, n in enumerate(set(meta_src) | set(meta_dst))}
+        mapped_meta_src = [ntypes_invmap[v] for v in meta_src]
+        mapped_meta_dst = [ntypes_invmap[v] for v in meta_dst]
+        node_frames = [self._node_frames[i] for i in ntypes_invmap]
         edge_frames = [self._edge_frames[i] for i in etype_ids]
-        induced_ntypes = [self._ntypes[i] for i in induced_ntype_ids]
+        induced_ntypes = [self._ntypes[i] for i in ntypes_invmap]
         induced_etypes = [self._etypes[i] for i in etype_ids]   # get the "name" of edge type
 
         metagraph = graph_index.from_edge_list((mapped_meta_src, mapped_meta_dst), True, True)


Test patch:
diff --git a/tests/compute/test_heterograph.py b/tests/compute/test_heterograph.py
--- a/tests/compute/test_heterograph.py
+++ b/tests/compute/test_heterograph.py
@@ -732,7 +732,7 @@ def _check_subgraph(g, sg):
     sg2 = g.edge_subgraph({'follows': [1], 'plays': [1], 'wishes': [1]})
     _check_subgraph(g, sg2)
 
-    def _check_typed_subgraph(g, sg):
+    def _check_typed_subgraph1(g, sg):
         assert set(sg.ntypes) == {'user', 'game'}
         assert set(sg.etypes) == {'follows', 'plays', 'wishes'}
         for ntype in sg.ntypes:
@@ -749,10 +749,23 @@ def _check_typed_subgraph(g, sg):
         assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'])
         assert F.array_equal(sg.edges['follows'].data['h'], g.edges['follows'].data['h'])
 
+    def _check_typed_subgraph2(g, sg):
+        assert set(sg.ntypes) == {'developer', 'game'}
+        assert set(sg.etypes) == {'develops'}
+        for ntype in sg.ntypes:
+            assert sg.number_of_nodes(ntype) == g.number_of_nodes(ntype)
+        for etype in sg.etypes:
+            src_sg, dst_sg = sg.all_edges(etype=etype, order='eid')
+            src_g, dst_g = g.all_edges(etype=etype, order='eid')
+            assert F.array_equal(src_sg, src_g)
+            assert F.array_equal(dst_sg, dst_g)
+
     sg3 = g.node_type_subgraph(['user', 'game'])
-    _check_typed_subgraph(g, sg3)
-    sg4 = g.edge_type_subgraph(['follows', 'plays', 'wishes'])
-    _check_typed_subgraph(g, sg4)
+    _check_typed_subgraph1(g, sg3)
+    sg4 = g.edge_type_subgraph(['develops'])
+    _check_typed_subgraph2(g, sg4)
+    sg5 = g.edge_type_subgraph(['follows', 'plays', 'wishes'])
+    _check_typed_subgraph1(g, sg5)
 
 def test_apply():
     def node_udf(nodes):