new: lazy attempt at #4

arangoml · Aug 3, 2022 · aca385a · aca385a
1 parent 4bb66f1
commit aca385a
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/adbpyg_adapter/encoders.py b/adbpyg_adapter/encoders.py
@@ -5,7 +5,12 @@
 
 from typing import Any, Dict, Optional
 
-from pandas import DataFrame
+try:
+    # https://github.com/arangoml/pyg-adapter/issues/4
+    from cudf import DataFrame
+except ModuleNotFoundError:
+    from pandas import DataFrame
+
 from torch import Tensor, from_numpy, zeros
 
 

diff --git a/adbpyg_adapter/typings.py b/adbpyg_adapter/typings.py
@@ -10,7 +10,12 @@
 
 from typing import Any, Callable, DefaultDict, Dict, List, Tuple, Union
 
-from pandas import DataFrame
+try:
+    # https://github.com/arangoml/pyg-adapter/issues/4
+    from cudf import DataFrame
+except ModuleNotFoundError:
+    from pandas import DataFrame
+
 from torch import Tensor
 
 Json = Dict[str, Any]

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -4,10 +4,15 @@
 from pathlib import Path
 from typing import Any, Callable
 
+try:
+    # https://github.com/arangoml/pyg-adapter/issues/4
+    from cudf import DataFrame
+except ModuleNotFoundError:
+    from pandas import DataFrame
+
 from arango import ArangoClient
 from arango.database import StandardDatabase
 from arango.http import DefaultHTTPClient
-from pandas import DataFrame
 from torch import Tensor, tensor
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.datasets import Amazon, FakeDataset, FakeHeteroDataset, KarateClub