[BugFix] skip frames whose num_rows is zero in dgl.heterograph.combine_frames() (#3110)

Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>

python/dgl/heterograph.py

@@ -6081,20 +6081,23 @@ def combine_frames(frames, ids, col_names=None):
         The resulting frame
     """
     # find common columns and check if their schemes match
-    if col_names is None:
-        schemes = {key: scheme for key, scheme in frames[ids[0]].schemes.items()}
-    else:
-        schemes = {key: frames[ids[0]].schemes[key] for key in col_names}
+    schemes = None
     for frame_id in ids:
         frame = frames[frame_id]
-        if frame.num_rows != 0:
-            for key, scheme in list(schemes.items()):
-                if key in frame.schemes:
-                    if frame.schemes[key] != scheme:
-                        raise DGLError('Cannot concatenate column %s with shape %s and shape %s' %
-                                       (key, frame.schemes[key], scheme))
-                else:
-                    del schemes[key]
+        if frame.num_rows == 0:
+            continue
+        if schemes is None:
+            schemes = frame.schemes
+            if col_names is not None:
+                schemes = {key: frame.schemes[key] for key in col_names}
+            continue
+        for key, scheme in list(schemes.items()):
+            if key in frame.schemes:
+                if frame.schemes[key] != scheme:
+                    raise DGLError('Cannot concatenate column %s with shape %s and shape %s' %
+                                   (key, frame.schemes[key], scheme))
+            else:
+                del schemes[key]
 
     if len(schemes) == 0:
         return None

tests/compute/test_heterograph.py

@@ -1082,6 +1082,22 @@ def test_convert(idtype):
     assert hg.device == g.device
     assert g.number_of_nodes() == 5
 
+    # hetero_to_subgraph_to_homo
+    hg = dgl.heterograph({
+        ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 2, 1]),
+        ('user', 'follows', 'user'): ([0, 1, 1], [1, 2, 2])
+    }, idtype=idtype, device=F.ctx())
+    hg.nodes['user'].data['h'] = F.copy_to(
+        F.tensor([[1, 0], [0, 1], [1, 1]], dtype=idtype), ctx=F.ctx())
+    sg = dgl.node_subgraph(hg, {'user': [1, 2]})
+    assert len(sg.ntypes) == 2
+    assert len(sg.etypes) == 2
+    assert sg.num_nodes('user') == 2
+    assert sg.num_nodes('game') == 0
+    g = dgl.to_homogeneous(sg, ndata=['h'])
+    assert 'h' in g.ndata.keys()
+    assert g.num_nodes() == 2
+
 @unittest.skipIf(F._default_context_str == 'gpu', reason="Test on cpu is enough")
 @parametrize_dtype
 def test_to_homo_zero_nodes(idtype):