graphdatascience 1.19__tar.gz → 1.20a1__tar.gz

{graphdatascience-1.19/src/graphdatascience.egg-info → graphdatascience-1.20a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.19
+Version: 1.20a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Author-email: Neo4j <team-gds@neo4j.org>
 License-Expression: Apache-2.0
@@ -18,6 +18,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Database
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development
@@ -29,7 +30,7 @@ Requires-Dist: multimethod<3.0,>=1.0
 Requires-Dist: neo4j<7.0,>=4.4.12
 Requires-Dist: numpy<2.4
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<23.0,>=17.0
+Requires-Dist: pyarrow<24.0,>=18.0
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
 Requires-Dist: typing-extensions<5.0,>=4.0

{graphdatascience-1.19 → graphdatascience-1.20a1}/pyproject.toml

@@ -20,6 +20,7 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
     "Topic :: Database",
     "Topic :: Scientific/Engineering",
     "Topic :: Software Development",
@@ -31,7 +32,7 @@ dependencies = [
     "neo4j >= 4.4.12, < 7.0",
     "numpy < 2.4",
     "pandas >= 1.0, < 3.0",
-    "pyarrow >= 17.0, < 23.0",
+    "pyarrow >= 18.0, < 24.0",
     "textdistance >= 4.0, < 5.0",
     "tqdm >= 4.0, < 5.0",
     "typing-extensions >= 4.0, < 5.0",
@@ -67,8 +68,8 @@ dev-base = [
     "mypy == 1.19.1",
     "pandas-stubs == 2.2.3.241009",
     "tox == 4.34.1",
-    "types-setuptools == 75.8.0.20250110",
-    "enum-tools[sphinx] == 0.12.0",
+    "types-setuptools == 80.9.0.20251223",
+    "enum-tools[sphinx]==0.13.0",
     "autodoc_pydantic",
     "types-requests",
     "types-tqdm",
@@ -86,7 +87,7 @@ test = [
 
 docs-ci = [
     "sphinx == 7.3.7",
-    "enum-tools[sphinx] == 0.12.0",
+    "enum-tools[sphinx]==0.13.0",
     "autodoc_pydantic",
 ]
 notebook-base = [
@@ -95,7 +96,10 @@ notebook-base = [
     "nbclient==0.10.4",
     "ipykernel==7.1.0",
 ]
-notebook-aura-ci = [{ include-group = "notebook-base" }, "python-dotenv==1.2.1"]
+notebook-aura-ci = [
+    { include-group = "notebook-base" },
+    "python-dotenv==1.2.1"
+]
 notebook-ci = [
     { include-group = "notebook-base" },
     "scipy == 1.14.0",
@@ -107,7 +111,6 @@ notebook-ci = [
 
 publish = ["build", "twine"]
 
-
 [tool.setuptools]
 include-package-data = true
 zip-safe = false

graphdatascience-1.20a1/src/graphdatascience/arrow_client/progress_callback.py

@@ -0,0 +1,5 @@
+from typing import Protocol
+
+
+class ProgressCallback(Protocol):
+    def __call__(self, num_rows: int) -> None: ...

{graphdatascience-1.19 → graphdatascience-1.20a1}/src/graphdatascience/arrow_client/v1/gds_arrow_client.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 import logging
 from types import TracebackType
-from typing import Any, Callable, Iterable, Type
+from typing import Any, Iterable, Type
 
 import pandas
 import pyarrow
@@ -17,6 +17,7 @@ from graphdatascience.arrow_client.v1.data_mapper_utils import deserialize_singl
 
 from ...procedure_surface.arrow.error_handler import handle_flight_error
 from ...semantic_version.semantic_version import SemanticVersion
+from ..progress_callback import ProgressCallback
 
 
 class GdsArrowClient:
@@ -386,7 +387,7 @@ class GdsArrowClient:
         graph_name: str,
         node_data: pyarrow.Table | Iterable[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10_000,
-        progress_callback: Callable[[int], None] = lambda x: None,
+        progress_callback: ProgressCallback = lambda x: None,
     ) -> None:
         """
         Uploads node data to the server.
@@ -399,7 +400,7 @@ class GdsArrowClient:
             The node data to upload
         batch_size : int
             The number of rows per batch
-        progress_callback : Callable[[int], None]
+        progress_callback : ProgressCallback
             A callback function that is called with the number of rows uploaded after each batch
         """
         self._upload_data(graph_name, "node", node_data, batch_size, progress_callback)
@@ -409,20 +410,20 @@ class GdsArrowClient:
         graph_name: str,
         relationship_data: pyarrow.Table | Iterable[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10_000,
-        progress_callback: Callable[[int], None] = lambda x: None,
+        progress_callback: ProgressCallback = lambda num_rows: None,
     ) -> None:
         """
         Uploads relationship data to the server.
 
         Parameters
         ----------
-        graph_name : str
+        graph_name
             The name of the import process
-        relationship_data : pyarrow.Table | Iterable[pyarrow.RecordBatch] | DataFrame
+        relationship_data
             The relationship data to upload
-        batch_size : int
+        batch_size
             The number of rows per batch
-        progress_callback : Callable[[int], None]
+        progress_callback
             A callback function that is called with the number of rows uploaded after each batch
         """
         self._upload_data(graph_name, "relationship", relationship_data, batch_size, progress_callback)
@@ -432,20 +433,20 @@ class GdsArrowClient:
         graph_name: str,
         triplet_data: pyarrow.Table | Iterable[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10_000,
-        progress_callback: Callable[[int], None] = lambda x: None,
+        progress_callback: ProgressCallback = lambda num_triplets: None,
     ) -> None:
         """
         Uploads triplet data to the server.
 
         Parameters
         ----------
-        graph_name : str
+        graph_name
             The name of the import process
         triplet_data : pyarrow.Table | Iterable[pyarrow.RecordBatch] | DataFrame
             The triplet data to upload
-        batch_size : int
+        batch_size
             The number of rows per batch
-        progress_callback : Callable[[int], None]
+        progress_callback
             A callback function that is called with the number of rows uploaded after each batch
         """
         self._upload_data(graph_name, "triplet", triplet_data, batch_size, progress_callback)
@@ -484,7 +485,7 @@ class GdsArrowClient:
         entity_type: str,
         data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int,
-        progress_callback: Callable[[int], None],
+        progress_callback: ProgressCallback,
     ) -> None:
         match data:
             case pyarrow.Table():

{graphdatascience-1.19 → graphdatascience-1.20a1}/src/graphdatascience/arrow_client/v2/gds_arrow_client.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 import logging
 from types import TracebackType
-from typing import Any, Callable, Type
+from typing import Any, Type
 
 import pandas
 import pyarrow
@@ -15,6 +15,7 @@ from graphdatascience.query_runner.termination_flag import TerminationFlag
 
 from ...procedure_surface.api.default_values import ALL_TYPES
 from ...procedure_surface.utils.config_converter import ConfigConverter
+from ..progress_callback import ProgressCallback
 from .api_types import JobStatus
 from .job_client import JobClient
 
@@ -326,9 +327,9 @@ class GdsArrowClient:
     def upload_nodes(
         self,
         job_id: str,
-        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame | list[pandas.DataFrame],
         batch_size: int = 10000,
-        progress_callback: Callable[[int], None] = lambda x: None,
+        progress_callback: ProgressCallback = lambda num_rows: None,
         termination_flag: TerminationFlag | None = None,
     ) -> None:
         """
@@ -354,9 +355,9 @@ class GdsArrowClient:
     def upload_relationships(
         self,
         job_id: str,
-        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame | list[pandas.DataFrame],
         batch_size: int = 10000,
-        progress_callback: Callable[[int], None] = lambda x: None,
+        progress_callback: ProgressCallback = lambda num_rows: None,
         termination_flag: TerminationFlag | None = None,
     ) -> None:
         """
@@ -384,7 +385,7 @@ class GdsArrowClient:
         job_id: str,
         data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10000,
-        progress_callback: Callable[[int], None] = lambda x: None,
+        progress_callback: ProgressCallback = lambda num_triplets: None,
         termination_flag: TerminationFlag | None = None,
     ) -> None:
         """
@@ -475,18 +476,27 @@ class GdsArrowClient:
         self,
         endpoint: str,
         job_id: str,
-        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | list[pandas.DataFrame],
         batch_size: int = 10000,
-        progress_callback: Callable[[int], None] = lambda x: None,
+        progress_callback: ProgressCallback = lambda num_rows: None,
         termination_flag: TerminationFlag | None = None,
     ) -> None:
+        if isinstance(data, pandas.DataFrame):
+            data = [data]
+
+        def map_list_entry_to_batches(e: pyarrow.RecordBatch | pandas.DataFrame) -> list[pyarrow.RecordBatch]:
+            return [e] if isinstance(e, pyarrow.RecordBatch) else pyarrow.Table.from_pandas(e).to_batches(batch_size)
+
+        batches: list[pyarrow.RecordBatch] = []
         match data:
+            case list():
+                for entry in data:
+                    partial_batches = map_list_entry_to_batches(entry)
+                    batches.extend(partial_batches)
             case pyarrow.Table():
                 batches = data.to_batches(batch_size)
-            case pandas.DataFrame():
-                batches = pyarrow.Table.from_pandas(data).to_batches(batch_size)
             case _:
-                batches = data
+                raise ValueError(f"Unsupported data type: {type(data)}")
 
         flight_descriptor = {
             "name": endpoint,

graphdatascience-1.20a1/src/graphdatascience/datasets/graph_constructor_func.py

@@ -0,0 +1,16 @@
+from typing import Protocol
+
+from pandas import DataFrame
+
+from graphdatascience.graph.v2 import GraphV2
+
+
+class GraphConstructorFunc(Protocol):
+    def __call__(
+        self,
+        graph_name: str,
+        nodes: DataFrame | list[DataFrame],
+        relationships: DataFrame | list[DataFrame] | None = None,
+        concurrency: int | None = None,
+        undirected_relationship_types: list[str] | None = None,
+    ) -> GraphV2: ...

{graphdatascience-1.19/src/graphdatascience/graph → graphdatascience-1.20a1/src/graphdatascience/datasets}/nx_loader.py

@@ -4,28 +4,37 @@ from typing import Any
 import networkx as nx
 from pandas import DataFrame
 
-from ..error.client_only_endpoint import client_only_endpoint
-from ..error.illegal_attr_checker import IllegalAttrChecker
-from ..error.uncallable_namespace import UncallableNamespace
-from ..server_version.compatible_with import compatible_with
-from ..server_version.server_version import ServerVersion
-from .graph_object import Graph
-
-
-class NXLoader(UncallableNamespace, IllegalAttrChecker):
-    @client_only_endpoint("gds.graph.networkx")
-    @compatible_with("load", min_inclusive=ServerVersion(2, 1, 0))
-    def load(self, nx_G: nx.Graph, graph_name: str, concurrency: int = 4) -> Graph:
+from graphdatascience.datasets.graph_constructor_func import GraphConstructorFunc
+from graphdatascience.graph.v2 import GraphV2
+
+
+class NXLoader:
+    def __init__(self, graph_constructor: GraphConstructorFunc) -> None:
+        self._construct = graph_constructor
+
+    def load(self, nx_G: nx.Graph, graph_name: str, concurrency: int | None = None) -> GraphV2:
+        """
+        Parameters
+        ----------
+        nx_G
+            The networkx graph to load.
+        graph_name
+            Name of the graph to create.
+        concurrency
+            Number of concurrent threads to use.
+
+        Returns
+        -------
+        GraphV2
+            A handle to the graph.
+        """
         nodes, rels = self._parse(nx_G)
 
         undirected_rel_types = []
         if not isinstance(nx_G, nx_G.to_directed_class()):
             undirected_rel_types = [df["relationshipType"][0] for df in rels]
 
-        constructor = self._query_runner.create_graph_constructor(graph_name, concurrency, undirected_rel_types)
-        constructor.run(nodes, rels)
-
-        return Graph(graph_name, self._query_runner)
+        return self._construct(graph_name, nodes, rels, concurrency, undirected_rel_types)
 
     @staticmethod
     def _attr_to_labels_key(labels_attr: Any, node_id: Any, no_node_labels: bool | None) -> tuple[str, ...]:

{graphdatascience-1.19/src/graphdatascience/graph → graphdatascience-1.20a1/src/graphdatascience/datasets}/ogb_loader.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from abc import ABC, abstractmethod
 from typing import Any, Protocol, TypedDict
@@ -7,13 +9,8 @@ import numpy as np
 import numpy.typing as npt
 import pandas as pd
 
-from ..error.client_only_endpoint import client_only_endpoint
-from ..error.illegal_attr_checker import IllegalAttrChecker
-from ..error.uncallable_namespace import UncallableNamespace
-from ..graph.graph_object import Graph
-from ..query_runner.query_runner import QueryRunner
-from ..server_version.compatible_with import compatible_with
-from ..server_version.server_version import ServerVersion
+from graphdatascience.datasets.graph_constructor_func import GraphConstructorFunc
+from graphdatascience.graph.v2 import GraphV2
 
 
 class _HomogeneousOGBGraphBase(TypedDict):
@@ -41,7 +38,7 @@ class HomogeneousOGBNDataset(Protocol):
     # `labels` here refers to class labels, not node labels in the Neo4j sense
     # The representation is a node_count x 1 shaped matrix
     labels: npt.NDArray[np.int64]
-    meta_info: "pd.Series[Any]"
+    meta_info: pd.Series[Any]
 
     @abstractmethod
     def get_idx_split(self) -> dict[str, npt.NDArray[np.int64]]:
@@ -50,7 +47,7 @@ class HomogeneousOGBNDataset(Protocol):
 
 class HomogeneousOGBLDataset(Protocol):
     graph: HomogeneousOGBGraph
-    meta_info: "pd.Series[Any]"
+    meta_info: pd.Series[Any]
     name: str
 
     @abstractmethod
@@ -63,7 +60,7 @@ class HeterogeneousOGBNDataset(Protocol):
     # `labels` here refers to class labels, not node labels in the Neo4j sense
     # The representation is a node_count x 1 shaped matrix
     labels: dict[str, npt.NDArray[np.int64]]
-    meta_info: "pd.Series[Any]"
+    meta_info: pd.Series[Any]
 
     @abstractmethod
     def get_idx_split(self) -> dict[str, dict[str, npt.NDArray[np.int64]]]:
@@ -72,26 +69,18 @@ class HeterogeneousOGBNDataset(Protocol):
 
 class HeterogeneousOGBLDataset(Protocol):
     graph: HeterogeneousOGBGraph
-    meta_info: "pd.Series[Any]"
+    meta_info: pd.Series[Any]
 
     @abstractmethod
     def get_edge_split(self) -> dict[str, dict[str, Any]]:
         pass
 
 
-class OGBLoader(UncallableNamespace, IllegalAttrChecker, ABC):
-    def __init__(self, query_runner: QueryRunner, namespace: str, server_version: ServerVersion):
-        self._query_runner = query_runner
-        self._namespace = namespace
-        self._server_version = server_version
+class OGBLoader(ABC):
+    def __init__(self, graph_constructor: GraphConstructorFunc):
+        self._construct = graph_constructor
         self._logger = logging.getLogger()
 
-    def _load(self, graph_name: str, nodes: list[pd.DataFrame], rels: list[pd.DataFrame], concurrency: int) -> Graph:
-        constructor = self._query_runner.create_graph_constructor(graph_name, concurrency, [])
-        constructor.run(nodes, rels)
-
-        return Graph(graph_name, self._query_runner)
-
 
 class OGBNLoader(OGBLoader):
     def _parse_homogeneous(self, dataset: HomogeneousOGBNDataset) -> tuple[list[pd.DataFrame], list[pd.DataFrame]]:
@@ -199,15 +188,30 @@ class OGBNLoader(OGBLoader):
 
         return nodes, rels
 
-    @client_only_endpoint("gds.graph.ogbn")
-    @compatible_with("load", min_inclusive=ServerVersion(2, 1, 0))
     def load(
         self,
         dataset_name: str,
         dataset_root_path: str = "dataset",
         graph_name: str | None = None,
-        concurrency: int = 4,
-    ) -> Graph:
+        concurrency: int | None = None,
+    ) -> GraphV2:
+        """
+        Parameters
+        ----------
+        dataset_name
+            Name of the dataset to load.
+        dataset_root_path
+            Path to the root directory of the dataset.
+        graph_name
+            Name of the graph to create.
+        concurrency
+            Number of concurrent threads to use.
+
+        Returns
+        -------
+        GraphV2
+            A handle to the graph.
+        """
         try:
             from ogb.nodeproppred import NodePropPredDataset
         except ModuleNotFoundError:
@@ -226,7 +230,7 @@ class OGBNLoader(OGBLoader):
         if not graph_name:
             graph_name = dataset_name
 
-        return self._load(graph_name, nodes, rels, concurrency)
+        return self._construct(graph_name, nodes, rels, concurrency=concurrency)
 
 
 class OGBLLoader(OGBLoader):
@@ -247,17 +251,7 @@ class OGBLLoader(OGBLoader):
         nodes = pd.DataFrame(node_dict)
 
         self._logger.info("Preparing relationship data for transfer to server...")
-
-        source_ids: list[int] = []
-        target_ids: list[int] = []
-        rel_types: list[str] = []
-        split = dataset.get_edge_split()
-
-        self._load_homogenous_ogbl_relationships(dataset.name, split, source_ids, target_ids, rel_types)
-
-        relationships = pd.DataFrame(
-            {"sourceNodeId": source_ids, "targetNodeId": target_ids, "relationshipType": rel_types}
-        )
+        relationships = self._load_homogenous_ogbl_relationships(dataset.name, dataset.get_edge_split())
 
         return [nodes], [relationships]
 
@@ -326,15 +320,30 @@ class OGBLLoader(OGBLoader):
 
         return nodes, rels
 
-    @client_only_endpoint("gds.graph.ogbl")
-    @compatible_with("load", min_inclusive=ServerVersion(2, 1, 0))
     def load(
         self,
         dataset_name: str,
         dataset_root_path: str = "dataset",
         graph_name: str | None = None,
-        concurrency: int = 4,
-    ) -> Graph:
+        concurrency: int | None = None,
+    ) -> GraphV2:
+        """
+        Parameters
+        ----------
+        dataset_name
+            Name of the dataset to load.
+        dataset_root_path
+            Path to the root directory of the dataset.
+        graph_name
+            Name of the graph to create.
+        concurrency
+            Number of concurrent threads to use.
+
+        Returns
+        -------
+        GraphV2
+            A handle to the graph.
+        """
         try:
             from ogb.linkproppred import LinkPropPredDataset
         except ModuleNotFoundError:
@@ -353,16 +362,16 @@ class OGBLLoader(OGBLoader):
         if not graph_name:
             graph_name = dataset_name
 
-        return self._load(graph_name, nodes, rels, concurrency)
+        return self._construct(graph_name, nodes, rels, concurrency=concurrency)
 
     def _load_homogenous_ogbl_relationships(
         self,
         dataset_name: str,
         split: dict[str, Any],
-        source_ids: list[int],
-        target_ids: list[int],
-        rel_types: list[str],
-    ) -> None:
+    ) -> pd.DataFrame:
+        source_ids: list[int] = []
+        target_ids: list[int] = []
+        rel_types: list[str] = []
         if dataset_name == "ogbl-wikikg2":
             for set_type, entity in split.items():
                 rel_suffix = f"{set_type.upper()}"
@@ -387,3 +396,5 @@ class OGBLLoader(OGBLoader):
                         source_ids.append(source_id)
                         target_ids.append(target_id)
                         rel_types.append(rel_type)
+
+        return pd.DataFrame({"sourceNodeId": source_ids, "targetNodeId": target_ids, "relationshipType": rel_types})

graphdatascience-1.20a1/src/graphdatascience/datasets/simple_file_loader.py

@@ -0,0 +1,79 @@
+import pathlib
+from typing import NamedTuple
+
+from neo4j import __version__ as neo4j_driver_version
+from pandas import DataFrame, read_parquet
+
+
+class GraphResources(NamedTuple):
+    nodes: list[DataFrame]
+    rels: list[DataFrame]
+
+
+class SimpleDatasetLoader:
+    def __init__(self) -> None:
+        self._is_neo4j_4_driver = neo4j_driver_version.startswith("4.")
+
+    @staticmethod
+    def _path(package: str, resource: str) -> pathlib.Path:
+        from importlib.resources import files
+
+        # files() returns a Traversable, but usages require a Path object
+        return pathlib.Path(str(files(package) / resource))
+
+    def cora(self) -> GraphResources:
+        file = self._path("graphdatascience.resources.cora", "cora_nodes.parquet.gzip")
+        nodes = read_parquet(file)
+
+        if self._is_neo4j_4_driver:
+            # features is read as an ndarray which was not supported in neo4j 4
+            nodes["features"] = nodes["features"].apply(lambda x: x.tolist())
+
+        rels = read_parquet(self._path("graphdatascience.resources.cora", "cora_rels.parquet.gzip"))
+        return GraphResources([nodes], [rels])
+
+    def karate_club(self) -> GraphResources:
+        nodes = DataFrame({"nodeId": range(1, 35)})
+        nodes["labels"] = "Person"
+
+        rels = read_parquet(self._path("graphdatascience.resources.karate", "karate_club.parquet.gzip"))
+
+        return GraphResources([nodes], [rels])
+
+    def imdb(self) -> GraphResources:
+        package = "graphdatascience.resources.imdb"
+        nodes = ["movies_with_genre", "movies_without_genre", "actors", "directors"]
+        rels = ["acted_in", "directed_in"]
+
+        node_dfs = []
+        for n in nodes:
+            resource = self._path(package, f"imdb_{n}.parquet.gzip")
+            df = read_parquet(resource)
+            if self._is_neo4j_4_driver:
+                # features is read as an ndarray which was not supported in neo4j 4
+                df["plot_keywords"] = df["plot_keywords"].apply(lambda x: x.tolist())
+            node_dfs.append(df)
+
+        rel_dfs = []
+        for r in rels:
+            resource = self._path(package, f"imdb_{r}.parquet.gzip")
+            rel_dfs.append(read_parquet(resource))
+        return GraphResources(node_dfs, rel_dfs)
+
+    def lastfm(self) -> GraphResources:
+        nodes = ["user_nodes", "artist_nodes"]
+        rels = ["user_friend_df_directed", "user_listen_artist_rels", "user_tag_artist_rels"]
+
+        package = "graphdatascience.resources.lastfm"
+
+        node_dfs = []
+        for n in nodes:
+            resource = self._path(package, f"{n}.parquet.gzip")
+            node_dfs.append(read_parquet(resource))
+
+        rel_dfs = []
+        for r in rels:
+            resource = self._path(package, f"{r}.parquet.gzip")
+            rel_dfs.append(read_parquet(resource))
+
+        return GraphResources(node_dfs, rel_dfs)

graphdatascience-1.20a1/src/graphdatascience/graph/nx_loader.py

@@ -0,0 +1,46 @@
+import networkx as nx
+import pandas as pd
+
+from graphdatascience.datasets.nx_loader import NXLoader as NXLoaderV2
+
+from ..error.client_only_endpoint import client_only_endpoint
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..error.uncallable_namespace import UncallableNamespace
+from ..server_version.compatible_with import compatible_with
+from ..server_version.server_version import ServerVersion
+from .graph_object import Graph
+from .v2 import GraphV2
+from .v2.graph_backend_cypher import get_graph
+
+
+class NXLoader(UncallableNamespace, IllegalAttrChecker):
+    def _construct_graph(
+        self,
+        graph_name: str,
+        nodes: pd.DataFrame | list[pd.DataFrame],
+        relationships: pd.DataFrame | list[pd.DataFrame] | None = None,
+        concurrency: int | None = None,
+        undirected_relationship_types: list[str] | None = None,
+    ) -> GraphV2:
+        if concurrency is None:
+            concurrency = 4
+
+        constructor = self._query_runner.create_graph_constructor(
+            graph_name, concurrency, undirected_relationship_types
+        )
+
+        if isinstance(nodes, pd.DataFrame):
+            nodes = [nodes]
+        if isinstance(relationships, pd.DataFrame):
+            relationships = [relationships]
+        if relationships is None:
+            relationships = []
+
+        constructor.run(nodes, relationships)
+        return get_graph(graph_name, self._query_runner)
+
+    @client_only_endpoint("gds.graph.networkx")
+    @compatible_with("load", min_inclusive=ServerVersion(2, 1, 0))
+    def load(self, nx_G: nx.Graph, graph_name: str, concurrency: int = 4) -> Graph:
+        G_V2 = NXLoaderV2(self._construct_graph).load(nx_G, graph_name, concurrency)
+        return Graph(G_V2.name(), self._query_runner)