Fix PyG conversion issues

kuzudb · Feb 5, 2023 · b360738 · b360738
1 parent 8cc93cb
commit b360738
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/tools/python_api/src_py/torch_geometric_result_converter.py b/tools/python_api/src_py/torch_geometric_result_converter.py
@@ -93,10 +93,9 @@ def __extract_properties_from_node(self, node, label, node_property_names):
                 pos = self.__add_unconverted_property(node, label, prop_name)
                 continue
 
-            # Read primary key but do not add it to the node properties
+            # Read primary key
             if node_property_names[prop_name]["is_primary_key"]:
                 primary_key = node[prop_name]
-                continue
 
             # Mark properties that are not supported by torch_geometric as unconverted
             if node_property_names[prop_name]["type"] not in [Type.INT64.value, Type.DOUBLE.value, Type.BOOL.value]:
@@ -249,8 +248,13 @@ def __convert_to_torch_geometric(self):
         pos_to_primary_key_dict = self.pos_to_primary_key_dict[
             label] if not is_hetero else self.pos_to_primary_key_dict
 
-        unconverted_properties = self.unconverted_properties if is_hetero else self.unconverted_properties[next(
-            iter(self.unconverted_properties))]
+        if is_hetero:
+            unconverted_properties = self.unconverted_properties
+        elif len(self.unconverted_properties) == 0:
+            unconverted_properties = {}
+        else:
+            unconverted_properties = self.unconverted_properties[next(
+                iter(self.unconverted_properties))]
         return data, pos_to_primary_key_dict, unconverted_properties
 
     def get_as_torch_geometric(self):