fix: edge_index.dtype is float (#31)

* add `edge_index.dtype` assertions

(tests should fail)

* fix: `edge_index.dtype` to `int`

* fix: `edge_index` creation

* cleanup

adbpyg_adapter/adapter.py

@@ -5,6 +5,7 @@
 from math import ceil
 from typing import Any, Callable, DefaultDict, Dict, List, Optional, Set, Tuple, Union
 
+import numpy as np
 import torch
 from arango.cursor import Cursor
 from arango.database import StandardDatabase
@@ -904,24 +905,31 @@ def __process_adb_edge_df(
             ]
 
             # 5. Map each ArangoDB from/to _key to the corresponding PyG node id
-            from_nodes = et_df["from_key"].map(adb_map[from_col]).tolist()
-            to_nodes = et_df["to_key"].map(adb_map[to_col]).tolist()
+            # NOTE: map() is somehow converting int values to float...
+            # So we rely on astype(int) to convert the float back to int,
+            # but we also fill NaN values with -1 so that we can convert
+            # the entire column to int without any issues. Need to revisit...
+            from_n = et_df["from_key"].map(adb_map[from_col]).fillna(-1).astype(int)
+            to_n = et_df["to_key"].map(adb_map[to_col]).fillna(-1).astype(int)
 
             # 6. Set/Update the PyG Edge Index
-            edge_index = tensor([from_nodes, to_nodes])
-            edge_data.edge_index = torch.cat(
-                (edge_data.get("edge_index", tensor([])), edge_index), dim=1
+            edge_index = tensor(
+                np.array([from_n.to_numpy(), to_n.to_numpy()]), dtype=torch.int64
             )
 
+            empty_tensor = tensor([], dtype=torch.int64)
+            existing_edge_index = edge_data.get("edge_index", empty_tensor)
+            edge_data.edge_index = cat((existing_edge_index, edge_index), dim=1)
+
             # 7. Deal with invalid edges
-            if torch.any(torch.isnan(edge_data.edge_index)):
+            if torch.any(edge_data.edge_index == -1):
                 if strict:
                     m = f"Invalid edges found in Edge Collection {e_col}, {from_col} -> {to_col}."  # noqa: E501
                     raise InvalidADBEdgesError(m)
                 else:
                     # Remove the invalid edges
                     edge_data.edge_index = edge_data.edge_index[
-                        :, ~torch.any(edge_data.edge_index.isnan(), dim=0)
+                        :, ~torch.any(edge_data.edge_index == -1, dim=0)
                     ]
 
             # 8. Set the PyG Edge Data

tests/test_adapter.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, List, Optional, Set, Union
 
 import pytest
+import torch
 from pandas import DataFrame
 from torch import Tensor, cat, long, tensor
 from torch_geometric.data import Data, HeteroData
@@ -554,6 +555,7 @@ def test_adb_partial_to_pyg() -> None:
     assert type(pyg_g_new) is Data
     assert pyg_g["v0"].x.tolist() == pyg_g_new.x.tolist()
     assert pyg_g["v0"].y.tolist() == pyg_g_new.y.tolist()
+    assert pyg_g[e_t].edge_index.dtype == torch.int64
     assert pyg_g[e_t].edge_index.tolist() == pyg_g_new.edge_index.tolist()
     assert pyg_g[e_t].edge_attr.tolist() == pyg_g_new.edge_attr.tolist()
 
@@ -714,13 +716,17 @@ def test_adb_graph_to_pyg_to_arangodb_with_missing_document_and_permissive(
 
     graph = db.graph(name)
     v_cols: Set[str] = graph.vertex_collections()
+    assert len(v_cols) == 1
     edge_definitions: List[Json] = graph.edge_definitions()
     e_cols: Set[str] = {c["edge_collection"] for c in edge_definitions}
+    assert len(e_cols) == 1
 
     for v_col in v_cols:
         vertex_collection = db.collection(v_col)
         vertex_collection.delete("0")
 
+    number_of_missing_edges = 32  # (i.e node 0 has 32 edges)
+
     metagraph: ADBMetagraph = {
         "vertexCollections": {col: {} for col in v_cols},
         "edgeCollections": {col: {} for col in e_cols},
@@ -729,7 +735,8 @@ def test_adb_graph_to_pyg_to_arangodb_with_missing_document_and_permissive(
     data = adapter.arangodb_to_pyg(name, metagraph=metagraph, strict=False)
 
     collection_count: int = db.collection(list(e_cols)[0]).count()
-    assert len(data.edge_index[0]) < collection_count
+    assert data.edge_index.dtype == torch.int64
+    assert data.num_edges + number_of_missing_edges == collection_count
 
     db.delete_graph(name, drop_collections=True)
 
@@ -1076,6 +1083,7 @@ def assert_adb_to_pyg(
             from_nodes = et_df["from_key"].map(adb_map[from_col]).tolist()
             to_nodes = et_df["to_key"].map(adb_map[to_col]).tolist()
 
+            assert edge_data.edge_index.dtype == torch.int64
             assert from_nodes == edge_data.edge_index[0].tolist()
             assert to_nodes == edge_data.edge_index[1].tolist()