Merge pull request #574 from adamnsch/ogbl-heterogeneous-fix

Parse node IDs correctly for heterogeneous OGBL graphs

Author: adamnsch

changelog.md

@@ -12,10 +12,10 @@
 
 ## Bug fixes
 
+* Fixed an issue where source and target IDs of relationships in heterogeneous OGBL graphs were not parsed correctly.
 
 
 ## Improvements
 
 
-
 ## Other changes

graphdatascience/graph/ogb_loader.py

@@ -314,8 +314,8 @@ def _parse_heterogeneous(self, dataset: HeterogeneousOGBLDataset) -> Tuple[List[
                 assert source_labels[i] == source_label
                 assert target_labels[i] == target_label
 
-                source_ids[i] += node_id_offsets[edges["head_type"][i]]
-                target_ids[i] += node_id_offsets[edges["tail_type"][i]]
+                source_ids[i] += node_id_offsets[edges["head_type"][i]] + edges["head"][i]
+                target_ids[i] += node_id_offsets[edges["tail_type"][i]] + edges["tail"][i]
 
                 rel_types.append(f"{edge_type}_{set_type.upper()}")