new: adbdgl refactor

.github/workflows/build.yml

@@ -13,7 +13,7 @@ jobs:
  runs-on: ubuntu-latest
  strategy:
  matrix:
- python: ["3.7", "3.8", "3.9"]
+ python: ["3.8", "3.9", "3.10", "3.11"]
  name: Python ${{ matrix.python }}
  steps:
  - uses: actions/checkout@v2

.github/workflows/release.yml

@@ -11,7 +11,7 @@ jobs:
  runs-on: ubuntu-latest
  strategy:
  matrix:
- python: ["3.7", "3.8", "3.9"]
+ python: ["3.8", "3.9", "3.10", "3.11"]
  name: Python ${{ matrix.python }}
  steps:
  - uses: actions/checkout@v2

README.md

@@ -10,7 +10,7 @@
 
 [![License](https://img.shields.io/github/license/arangoml/dgl-adapter?color=9E2165&style=for-the-badge)](https://github.com/arangoml/dgl-adapter/blob/master/LICENSE)
 [![Code style: black](https://img.shields.io/static/v1?style=for-the-badge&label=code%20style&message=black&color=black)](https://github.com/psf/black)
-[![Downloads](https://img.shields.io/badge/dynamic/json?style=for-the-badge&color=282661&label=Downloads&query=total_downloads&url=https://api.pepy.tech/api/projects/adbdgl-adapter)](https://pepy.tech/project/adbdgl-adapter)
+[![Downloads](https://img.shields.io/badge/dynamic/json?style=for-the-badge&color=282661&label=Downloads&query=total_downloads&url=https://api.pepy.tech/api/v2/projects/adbdgl-adapter)](https://pepy.tech/project/adbdgl-adapter)
 
 
 <a href="https://www.arangodb.com/" rel="arangodb.com">![](https://raw.githubusercontent.com/arangoml/dgl-adapter/master/examples/assets/adb_logo.png)</a>
@@ -45,44 +45,186 @@ pip install git+https://github.com/arangoml/dgl-adapter.git
 Also available as an ArangoDB Lunch & Learn session: [Graph & Beyond Course #2.8](https://www.arangodb.com/resources/lunch-sessions/graph-beyond-lunch-break-2-8-dgl-adapter/)
 
 ```py
+import pandas
+import torch
+import dgl
+
 from arango import ArangoClient # Python-Arango driver
-from dgl.data import KarateClubDataset # Sample graph from DGL
 
-from adbdgl_adapter import ADBDGL_Adapter
+from adbdgl_adapter import ADBDGL_Adapter, ADBDGL_Controller
+from adbdgl_adapter.encoders import IdentityEncoder, CategoricalEncoder
 
-# Let's assume that the ArangoDB "fraud detection" dataset is imported to this endpoint
+# Let's assume that the ArangoDB "IMDB" dataset is imported to this endpoint
 db = ArangoClient(hosts="http://localhost:8529").db("_system", username="root", password="")
 
+fake_hetero = dgl.heterograph({
+ ("user", "follows", "user"): (torch.tensor([0, 1]), torch.tensor([1, 2])),
+ ("user", "follows", "topic"): (torch.tensor([1, 1]), torch.tensor([1, 2])),
+ ("user", "plays", "game"): (torch.tensor([0, 3]), torch.tensor([3, 4])),
+})
+fake_hetero.nodes["user"].data["features"] = torch.tensor([21, 44, 16, 25])
+fake_hetero.nodes["user"].data["label"] = torch.tensor([1, 2, 0, 1])
+fake_hetero.nodes["game"].data["features"] = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1], [1, 1]])
+fake_hetero.edges[("user", "plays", "game")].data["features"] = torch.tensor([[6, 1], [1000, 0]])
+
 adbdgl_adapter = ADBDGL_Adapter(db)
+```
+
+### DGL to ArangoDB
+```py
+# 1.1: DGL to ArangoDB
+adb_g = adbdgl_adapter.dgl_to_arangodb("FakeHetero", fake_hetero)
 
-# Use Case 1.1: ArangoDB to DGL via Graph name
-dgl_fraud_graph = adbdgl_adapter.arangodb_graph_to_dgl("fraud-detection")
+# 1.2: DGL to ArangoDB with a (completely optional) metagraph for customized adapter behaviour
+def label_tensor_to_2_column_dataframe(dgl_tensor, adb_df):
+ """
+ A user-defined function to create two
+ ArangoDB attributes out of the 'user' label tensor
+
+ :param dgl_tensor: The DGL Tensor containing the data
+ :type dgl_tensor: torch.Tensor
+ :param adb_df: The ArangoDB DataFrame to populate, whose
+ size is preset to the length of **dgl_tensor**.
+ :type adb_df: pandas.DataFrame
+
+ NOTE: user-defined functions must return the modified **adb_df**
+ """
+ label_map = {0: "Class A", 1: "Class B", 2: "Class C"}
+
+ adb_df["label_num"] = dgl_tensor.tolist()
+ adb_df["label_str"] = adb_df["label_num"].map(label_map)
+
+ return adb_df
 
-# Use Case 1.2: ArangoDB to DGL via Collection names
-dgl_fraud_graph_2 = adbdgl_adapter.arangodb_collections_to_dgl(
- "fraud-detection",
- {"account", "Class", "customer"}, # Vertex collections
- {"accountHolder", "Relationship", "transaction"}, # Edge collections
-)
 
-# Use Case 1.3: ArangoDB to DGL via Metagraph
 metagraph = {
+ "nodeTypes": {
+ "user": {
+ "features": "user_age", # 1) you can specify a string value for attribute renaming
+ "label": label_tensor_to_2_column_dataframe, # 2) you can specify a function for user-defined handling, as long as the function returns a Pandas DataFrame
+ },
+ # 3) You can specify set of strings if you want to preserve the same DGL attribute names for the node/edge type
+ "game": {"features"} # this is equivalent to {"features": "features"}
+ },
+ "edgeTypes": {
+ ("user", "plays", "game"): {
+ # 4) you can specify a list of strings for tensor dissasembly (if you know the number of node/edge features in advance)
+ "features": ["hours_played", "is_satisfied_with_game"]
+ },
+ },
+}
+
+
+adb_g = adbdgl_adapter.dgl_to_arangodb("FakeHetero", fake_hetero, metagraph, explicit_metagraph=False)
+
+# 1.3: DGL to ArangoDB with the same (optional) metagraph, but with `explicit_metagraph=True`
+# With `explicit_metagraph=True`, the node & edge types omitted from the metagraph will NOT be converted to ArangoDB.
+# Only 'user', 'game', and ('user', 'plays', 'game') will be brought over (i.e 'topic', ('user', 'follows', 'user'), ... are ignored)
+adb_g = adbdgl_adapter.dgl_to_arangodb("FakeHetero", fake_hetero, metagraph, explicit_metagraph=True)
+
+# 1.4: DGL to ArangoDB with a Custom Controller (more user-defined behavior)
+class Custom_ADBDGL_Controller(ADBDGL_Controller):
+ def _prepare_dgl_node(self, dgl_node: dict, node_type: str) -> dict:
+ """Optionally modify a DGL node object before it gets inserted into its designated ArangoDB collection.
+
+ :param dgl_node: The DGL node object to (optionally) modify.
+ :param node_type: The DGL Node Type of the node.
+ :return: The DGL Node object
+ """
+ dgl_node["foo"] = "bar"
+ return dgl_node
+
+ def _prepare_dgl_edge(self, dgl_edge: dict, edge_type: tuple) -> dict:
+ """Optionally modify a DGL edge object before it gets inserted into its designated ArangoDB collection.
+
+ :param dgl_edge: The DGL edge object to (optionally) modify.
+ :param edge_type: The Edge Type of the DGL edge. Formatted
+ as (from_collection, edge_collection, to_collection)
+ :return: The DGL Edge object
+ """
+ dgl_edge["bar"] = "foo"
+ return dgl_edge
+
+
+adb_g = ADBDGL_Adapter(db, Custom_ADBDGL_Controller()).dgl_to_arangodb("FakeHetero", fake_hetero)
+```
+
+### ArangoDB to DGL
+```py
+# Start from scratch!
+db.delete_graph("FakeHetero", drop_collections=True, ignore_missing=True)
+adbdgl_adapter.dgl_to_arangodb("FakeHetero", fake_hetero)
+
+# 2.1: ArangoDB to DGL via Graph name (does not transfer attributes)
+dgl_g = adbdgl_adapter.arangodb_graph_to_dgl("FakeHetero")
+
+# 2.2: ArangoDB to DGL via Collection names (does not transfer attributes)
+dgl_g = adbdgl_adapter.arangodb_collections_to_dgl("FakeHetero", v_cols={"user", "game"}, e_cols={"plays"})
+
+# 2.3: ArangoDB to DGL via Metagraph v1 (transfer attributes "as is", meaning they are already formatted to DGL data standards)
+# Learn more about the DGL Data Standards here: https://docs.dgl.ai/guide/graph.html#guide-graph
+metagraph_v1 = {
  "vertexCollections": {
- "account": {"Balance", "rank"},
- "customer": {"rank"},
- "Class": {},
+ # Move the "features" & "label" ArangoDB attributes to DGL as "features" & "label" Tensors
+ "user": {"features", "label"}, # equivalent to {"features": "features", "label": "label"}
+ "game": {"dgl_game_features": "features"},
+ "topic": {},
  },
  "edgeCollections": {
- "transaction": {"transaction_amt", "sender_bank_id", "receiver_bank_id"},
- "accountHolder": {},
- "Relationship": {},
+ "plays": {"dgl_plays_features": "features"}, 
+ "follows": {}
+ },
+}
+dgl_g = adbdgl_adapter.arangodb_to_dgl("FakeHetero", metagraph_v1)
+
+# 2.4: ArangoDB to DGL via Metagraph v2 (transfer attributes via user-defined encoders)
+# For more info on user-defined encoders, see https://pytorch-geometric.readthedocs.io/en/latest/notes/load_csv.html
+metagraph_v2 = {
+ "vertexCollections": {
+ "Movies": {
+ "features": { # Build a feature matrix from the "Action" & "Drama" document attributes
+ "Action": IdentityEncoder(dtype=torch.long),
+ "Drama": IdentityEncoder(dtype=torch.long),
+ },
+ "label": "Comedy",
+ },
+ "Users": {
+ "features": {
+ "Gender": CategoricalEncoder(), # CategoricalEncoder(mapping={"M": 0, "F": 1}),
+ "Age": IdentityEncoder(dtype=torch.long),
+ }
+ },
  },
+ "edgeCollections": {"Ratings": {"weight": "Rating"}},
 }
-dgl_fraud_graph_3 = adbdgl_adapter.arangodb_to_dgl("fraud-detection", metagraph)
+dgl_g = adbdgl_adapter.arangodb_to_dgl("IMDB", metagraph_v2)
+
+# 2.5: ArangoDB to DGL via Metagraph v3 (transfer attributes via user-defined functions)
+def udf_user_features(user_df):
+ # process the user_df Pandas DataFrame to return a feature matrix in a tensor
+ # user_df["features"] = ...
+ return torch.tensor(user_df["features"].to_list())
+
+
+def udf_game_features(game_df):
+ # process the game_df Pandas DataFrame to return a feature matrix in a tensor
+ # game_df["features"] = ...
+ return torch.tensor(game_df["features"].to_list())
+
 
-# Use Case 2: DGL to ArangoDB
-dgl_karate_graph = KarateClubDataset()[0]
-adb_karate_graph = adbdgl_adapter.dgl_to_arangodb("Karate", dgl_karate_graph)
+metagraph_v3 = {
+ "vertexCollections": {
+ "user": {
+ "features": udf_user_features, # supports named functions
+ "label": lambda df: torch.tensor(df["label"].to_list()), # also supports lambda functions
+ },
+ "game": {"features": udf_game_features},
+ },
+ "edgeCollections": {
+ "plays": {"features": (lambda df: torch.tensor(df["features"].to_list()))},
+ },
+}
+dgl_g = adbdgl_adapter.arangodb_to_dgl("FakeHetero", metagraph_v3)
 ```
 
 ## Development & Testing

adbdgl_adapter/abc.py

@@ -2,22 +2,20 @@
 # -*- coding: utf-8 -*-
 
 from abc import ABC
-from typing import Any, List, Set, Union
+from typing import Any, Set, Union
 
 from arango.graph import Graph as ArangoDBGraph
-from dgl import DGLGraph
-from dgl.heterograph import DGLHeteroGraph
-from torch import Tensor
+from dgl import DGLGraph, DGLHeteroGraph
 
-from .typings import ArangoMetagraph, DGLCanonicalEType, Json
+from .typings import ADBMetagraph, DGLCanonicalEType, DGLMetagraph, Json
 
 
 class Abstract_ADBDGL_Adapter(ABC):
  def __init__(self) -> None:
  raise NotImplementedError # pragma: no cover
 
  def arangodb_to_dgl(
- self, name: str, metagraph: ArangoMetagraph, **query_options: Any
+ self, name: str, metagraph: ADBMetagraph, **query_options: Any
  ) -> DGLHeteroGraph:
  raise NotImplementedError # pragma: no cover
 
@@ -33,39 +31,17 @@ def dgl_to_arangodb(
  self,
  name: str,
  dgl_g: Union[DGLGraph, DGLHeteroGraph],
+ metagraph: DGLMetagraph = {},
+ explicit_metagraph: bool = True,
  overwrite_graph: bool = False,
  **import_options: Any,
  ) -> ArangoDBGraph:
  raise NotImplementedError # pragma: no cover
 
- def etypes_to_edefinitions(
- self, canonical_etypes: List[DGLCanonicalEType]
- ) -> List[Json]:
- raise NotImplementedError # pragma: no cover
-
- def __prepare_dgl_features(self) -> None:
- raise NotImplementedError # pragma: no cover
-
- def __insert_dgl_features(self) -> None:
- raise NotImplementedError # pragma: no cover
-
- def __prepare_adb_attributes(self) -> None:
- raise NotImplementedError # pragma: no cover
-
- def __fetch_adb_docs(self) -> None:
- raise NotImplementedError # pragma: no cover
-
- def __insert_adb_docs(self) -> None:
- raise NotImplementedError # pragma: no cover
-
- @property
- def DEFAULT_CANONICAL_ETYPE(self) -> List[DGLCanonicalEType]:
- return [("_N", "_E", "_N")]
-
 
 class Abstract_ADBDGL_Controller(ABC):
- def _adb_attribute_to_dgl_feature(self, key: str, col: str, val: Any) -> Any:
+ def _prepare_dgl_node(self, dgl_node: Json, node_type: str) -> Json:
  raise NotImplementedError # pragma: no cover
 
- def _dgl_feature_to_adb_attribute(self, key: str, col: str, val: Tensor) -> Any:
+ def _prepare_dgl_edge(self, dgl_edge: Json, edge_type: DGLCanonicalEType) -> Json:
  raise NotImplementedError # pragma: no cover