Merge pull request #17 from arangoml/feature/exceptions

Add Strict Parameter & Handle Invalid Edges

adbpyg_adapter/adapter.py

@@ -4,6 +4,7 @@
 from collections import defaultdict
 from typing import Any, DefaultDict, Dict, List, Set, Union
 
+import torch
 from arango.database import Database
 from arango.graph import Graph as ADBGraph
 from pandas import DataFrame, Series
@@ -14,7 +15,7 @@
 
 from .abc import Abstract_ADBPyG_Adapter
 from .controller import ADBPyG_Controller
-from .exceptions import ADBMetagraphError, PyGMetagraphError
+from .exceptions import ADBMetagraphError, InvalidADBEdgesError, PyGMetagraphError
 from .typings import (
     ADBMap,
     ADBMetagraph,
@@ -79,6 +80,7 @@ def arangodb_to_pyg(
         name: str,
         metagraph: ADBMetagraph,
         preserve_adb_keys: bool = False,
+        strict: bool = True,
         **query_options: Any,
     ) -> Union[Data, HeteroData]:
         """Create a PyG graph from ArangoDB data. DOES carry
@@ -127,6 +129,9 @@ def arangodb_to_pyg(
             ArangoDB graph is Heterogeneous, the ArangoDB keys will be preserved
             under `_key` in your PyG graph.
         :type preserve_adb_keys: bool
+        :param strict: Set fault tolerance when loading a graph from ArangoDB. If set
+            to false, this will ignore invalid edges (e.g. dangling/half edges).
+        :type strict: bool
         :param query_options: Keyword arguments to specify AQL query options when
             fetching documents from the ArangoDB instance. Full parameter list:
             https://docs.python-arango.com/en/main/specs.html#arango.aql.AQL.execute
@@ -296,6 +301,18 @@ def udf_v1_x(v1_df):
 
                 edge_data: EdgeStorage = data if is_homogeneous else data[edge_type]
                 edge_data.edge_index = tensor([from_nodes, to_nodes])
+
+                if torch.any(torch.isnan(edge_data.edge_index)):
+                    if strict:
+                        raise InvalidADBEdgesError(
+                            f"Invalid edges found in Edge Collection {e_col}, {from_col} -> {to_col}."  # noqa: E501
+                        )
+                    else:
+                        # Remove the invalid edges
+                        edge_data.edge_index = edge_data.edge_index[
+                            :, ~torch.any(edge_data.edge_index.isnan(), dim=0)
+                        ]
+
                 self.__set_pyg_data(meta, edge_data, et_df)
 
                 if preserve_adb_keys:
@@ -654,28 +671,28 @@ def ntypes_to_ocollections(
         return list(orphan_collections)
 
     def __fetch_adb_docs(
-        self, col: str, empty_meta: bool, query_options: Any
+        self, col: str, meta_is_empty: bool, query_options: Any
     ) -> DataFrame:
         """Fetches ArangoDB documents within a collection. Returns the
             documents in a DataFrame.
 
         :param col: The ArangoDB collection.
         :type col: str
-        :param empty_meta: Set to True if the metagraph specification
+        :param meta_is_empty: Set to True if the metagraph specification
             for **col** is empty.
-        :type empty_meta: bool
+        :type meta_is_empty: bool
         :param query_options: Keyword arguments to specify AQL query options
             when fetching documents from the ArangoDB instance.
         :type query_options: Any
         :return: A DataFrame representing the ArangoDB documents.
         :rtype: pandas.DataFrame
         """
-        # Only return the entire document if **empty_meta** is False
+        # Only return the entire document if **meta_is_empty** is False
         aql = f"""
             FOR doc IN @@col
                 RETURN {
                     "{ _key: doc._key, _from: doc._from, _to: doc._to }"
-                    if empty_meta
+                    if meta_is_empty
                     else "doc"
                 }
         """

adbpyg_adapter/exceptions.py

@@ -17,3 +17,16 @@ class ADBMetagraphError(ADBPyGValidationError):
 
 class PyGMetagraphError(ADBPyGValidationError):
     """Invalid PyG Metagraph value"""
+
+
+##################
+#   ADB -> PyG   #
+##################
+
+
+class ADBPyGImportError(ADBPyGError):
+    """Errors on import from ArangoDB to PyG"""
+
+
+class InvalidADBEdgesError(ADBPyGImportError):
+    """Invalid edges on import from ArangoDB to PyG"""

adbpyg_adapter/typings.py

@@ -8,7 +8,7 @@
     "PyGMap",
 ]
 
-from typing import Any, Callable, DefaultDict, Dict, List, Tuple, Union
+from typing import Any, Callable, DefaultDict, Dict, List, Set, Tuple, Union
 
 from pandas import DataFrame
 from torch import Tensor
@@ -20,7 +20,7 @@
 
 ADBEncoders = Dict[str, DataFrameToTensor]
 ADBMetagraphValues = Union[str, DataFrameToTensor, ADBEncoders]
-ADBMetagraph = Dict[str, Dict[str, Dict[str, ADBMetagraphValues]]]
+ADBMetagraph = Dict[str, Dict[str, Union[Set[str], Dict[str, ADBMetagraphValues]]]]
 
 PyGDataTypes = Union[str, Tuple[str, str, str]]
 PyGMetagraphValues = Union[str, List[str], TensorToDataFrame]

tests/test_adapter.py

@@ -10,7 +10,11 @@
 
 from adbpyg_adapter import ADBPyG_Adapter
 from adbpyg_adapter.encoders import CategoricalEncoder, IdentityEncoder
-from adbpyg_adapter.exceptions import ADBMetagraphError, PyGMetagraphError
+from adbpyg_adapter.exceptions import (
+    ADBMetagraphError,
+    InvalidADBEdgesError,
+    PyGMetagraphError,
+)
 from adbpyg_adapter.typings import (
     ADBMap,
     ADBMetagraph,
@@ -643,6 +647,68 @@ def test_adb_graph_to_pyg(
         db.delete_graph(name, drop_collections=True)
 
 
+@pytest.mark.parametrize("adapter", [adbpyg_adapter])
+def test_adb_graph_to_pyg_to_arangodb_with_missing_document_and_strict(
+    adapter: ADBPyG_Adapter,
+) -> None:
+    name = "Karate_3"
+    data = get_karate_graph()
+    db.delete_graph(name, drop_collections=True, ignore_missing=True)
+
+    ADBPyG_Adapter(db).pyg_to_arangodb(name, data)
+
+    graph = db.graph(name)
+    v_cols: Set[str] = graph.vertex_collections()  # type: ignore
+    edge_definitions: List[Json] = graph.edge_definitions()  # type: ignore
+    e_cols: Set[str] = {c["edge_collection"] for c in edge_definitions}
+
+    for v_col in v_cols:
+        vertex_collection = db.collection(v_col)
+        vertex_collection.delete("0")
+
+    metagraph: ADBMetagraph = {
+        "vertexCollections": {col: {} for col in v_cols},
+        "edgeCollections": {col: {} for col in e_cols},
+    }
+
+    with pytest.raises(InvalidADBEdgesError):
+        adapter.arangodb_to_pyg(name, metagraph=metagraph, strict=True)
+
+    db.delete_graph(name, drop_collections=True)
+
+
+@pytest.mark.parametrize("adapter", [adbpyg_adapter])
+def test_adb_graph_to_pyg_to_arangodb_with_missing_document_and_permissive(
+    adapter: ADBPyG_Adapter,
+) -> None:
+    name = "Karate_3"
+    data = get_karate_graph()
+    db.delete_graph(name, drop_collections=True, ignore_missing=True)
+
+    ADBPyG_Adapter(db).pyg_to_arangodb(name, data)
+
+    graph = db.graph(name)
+    v_cols: Set[str] = graph.vertex_collections()  # type: ignore
+    edge_definitions: List[Json] = graph.edge_definitions()  # type: ignore
+    e_cols: Set[str] = {c["edge_collection"] for c in edge_definitions}
+
+    for v_col in v_cols:
+        vertex_collection = db.collection(v_col)
+        vertex_collection.delete("0")
+
+    metagraph: ADBMetagraph = {
+        "vertexCollections": {col: {} for col in v_cols},
+        "edgeCollections": {col: {} for col in e_cols},
+    }
+
+    data = adapter.arangodb_to_pyg(name, metagraph=metagraph, strict=False)
+
+    collection_count: int = db.collection(list(e_cols)[0]).count()  # type: ignore
+    assert len(data.edge_index[0]) < collection_count
+
+    db.delete_graph(name, drop_collections=True)
+
+
 def test_full_cycle_imdb_without_preserve_adb_keys() -> None:
     name = "imdb"
     db.delete_graph(name, drop_collections=True, ignore_missing=True)
@@ -1000,7 +1066,7 @@ def assert_adb_to_pyg(
 
 
 def assert_adb_to_pyg_meta(
-    meta: Union[str, Dict[str, ADBMetagraphValues]],
+    meta: Union[str, Set[str], Dict[str, ADBMetagraphValues]],
     df: DataFrame,
     pyg_data: Union[NodeStorage, EdgeStorage],
 ) -> None: