phenolrs 0.5.12__cp313-cp313-manylinux_2_34_aarch64.whl

phenolrs/aql/typings.py

@@ -0,0 +1,98 @@
+"""Type definitions for AQL-based graph loading.
+
+This module follows the design specification from:
+https://github.com/arangodb/documents/blob/master/DesignDocuments/02_PLANNING/GetGraphsOutOfArangoDBviaAQL.md
+"""
+
+from typing import Any, Dict, List, Literal, TypedDict, Union
+
+# Supported data types for attributes
+DataType = Literal[
+    "bool", "string", "u64", "i64", "f64", "json", "number", "int", "float"
+]
+
+
+class AqlQuery(TypedDict, total=False):
+    """An AQL query with optional bind variables.
+
+    Each query should return items of the form:
+    {"vertices": [...], "edges": [...]}
+
+    Both vertices and edges are optional in the return value.
+    Vertices must have at least an _id attribute.
+    Edges must have at least _from and _to attributes.
+    """
+
+    query: str  # Required: The AQL query string
+    bindVars: Dict[str, Any]  # Optional: Bind variables for the query
+    bind_vars: Dict[str, Any]  # Alternative name for bindVars
+
+
+class AttributeItem(TypedDict):
+    """A single attribute definition with name and type."""
+
+    name: str
+    type: DataType
+
+
+# Attribute specification can be either:
+# - A dict mapping attribute names to types: {"name": "string", "age": "i64"}
+# - A list of attribute items: [{"name": "name", "type": "string"}, ...]
+AttributeSpec = Union[Dict[str, DataType], List[AttributeItem]]
+
+
+class DatabaseConfig(TypedDict, total=False):
+    """Database connection configuration."""
+
+    endpoints: List[str]  # List of ArangoDB endpoints
+    database: str  # Database name (default: "_system")
+    username: str  # Username for authentication
+    password: str  # Password for authentication
+    jwt_token: str  # JWT token for authentication (alternative to username/password)
+    tls_cert: str  # TLS certificate for secure connections
+
+
+class AqlDataLoadRequest(TypedDict, total=False):
+    """Request for AQL-based graph loading.
+
+    The queries field is a list of lists of AQL queries:
+    - The outer list is processed **sequentially**
+    - Each inner list of queries can be executed **in parallel**
+
+    This structure allows for:
+    1. Loading vertices first, then edges (for efficient memory usage)
+    2. Parallel loading of multiple vertex/edge collections
+    3. Sequential execution of dependent operations
+
+    Example for use case 1 (filtered collections)::
+
+        {
+            "queries": [
+                # First group: load all vertices in parallel
+                [
+                    {"query": "FOR x IN v1 RETURN {vertices: [x]}"},
+                    {"query": "FOR x IN v2 RETURN {vertices: [x]}"}
+                ],
+                # Second group: load all edges in parallel
+                [
+                    {"query": "FOR e IN edges1 RETURN {edges: [e]}"},
+                ]
+            ]
+        }
+
+    Example for use case 2 (graph traversals)::
+
+        {
+            "queries": [[{
+                "query": "FOR v, e IN 0..10 OUTBOUND @s GRAPH 'g' "
+                         "RETURN {vertices: [v], edges: [e]}",
+                "bindVars": {"s": "vertex/1"}
+            }]]
+        }
+    """
+
+    database_config: DatabaseConfig
+    batch_size: int  # Number of items per batch (default: 10000)
+    vertex_attributes: AttributeSpec  # Schema for vertex attributes
+    edge_attributes: AttributeSpec  # Schema for edge attributes
+    queries: List[List[AqlQuery]]  # List of query groups

phenolrs/networkx/__init__.py

@@ -0,0 +1 @@
+from .loader import NetworkXLoader  # noqa: F401

phenolrs/networkx/loader.py

@@ -0,0 +1,144 @@
+from typing import Any, Set, Tuple
+
+from phenolrs import PhenolError, graph_to_networkx_format
+
+from .typings import (
+    ArangoIDtoIndex,
+    DiGraphAdjDict,
+    DstIndices,
+    EdgeIndices,
+    EdgeValuesDict,
+    GraphAdjDict,
+    MultiDiGraphAdjDict,
+    MultiGraphAdjDict,
+    NodeDict,
+    SrcIndices,
+)
+
+
+class NetworkXLoader:
+    @staticmethod
+    def load_into_networkx(
+        database: str,
+        metagraph: dict[str, dict[str, Set[str]]],
+        hosts: list[str],
+        user_jwt: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        tls_cert: Any | None = None,
+        parallelism: int | None = None,
+        batch_size: int | None = None,
+        load_adj_dict: bool = True,
+        load_coo: bool = True,
+        load_all_vertex_attributes: bool = True,
+        load_all_edge_attributes: bool = True,
+        is_directed: bool = True,
+        is_multigraph: bool = True,
+        symmetrize_edges_if_directed: bool = False,
+    ) -> Tuple[
+        NodeDict,
+        GraphAdjDict | DiGraphAdjDict | MultiGraphAdjDict | MultiDiGraphAdjDict,
+        SrcIndices,
+        DstIndices,
+        EdgeIndices,
+        ArangoIDtoIndex,
+        EdgeValuesDict,
+    ]:
+        if "vertexCollections" not in metagraph:
+            raise PhenolError("vertexCollections not found in metagraph")
+
+        if "edgeCollections" not in metagraph:
+            raise PhenolError("edgeCollections not found in metagraph")
+
+        if len(metagraph["vertexCollections"]) + len(metagraph["edgeCollections"]) == 0:
+            m = "vertexCollections and edgeCollections cannot both be empty"
+            raise PhenolError(m)
+
+        if len(metagraph["edgeCollections"]) == 0 and (load_adj_dict or load_coo):
+            m = "edgeCollections must be non-empty if **load_adj_dict** or **load_coo** is True"  # noqa
+            raise PhenolError(m)
+
+        if load_all_vertex_attributes:
+            for entries in metagraph["vertexCollections"].values():
+                if len(entries) > 0:
+                    m = f"load_all_vertex_attributes is True, but a vertexCollections entry contains attributes: {entries}"  # noqa
+                    raise PhenolError(m)
+
+        if load_all_edge_attributes:
+            for entries in metagraph["edgeCollections"].values():
+                if len(entries) > 0:
+                    m = f"load_all_edge_attributes is True, but an edgeCollections entry contains attributes: {entries}"  # noqa
+                    raise PhenolError(m)
+
+        if len(metagraph["edgeCollections"]) != 0 and not (load_coo or load_adj_dict):
+            m = "load_coo and load_adj_dict cannot both be False if edgeCollections is non-empty"  # noqa
+            raise PhenolError(m)
+
+        # TODO: replace with pydantic validation
+        db_config_options: dict[str, Any] = {
+            "endpoints": hosts,
+            "database": database,
+        }
+
+        load_config_options: dict[str, Any] = {
+            "parallelism": parallelism if parallelism is not None else 8,
+            "batch_size": batch_size if batch_size is not None else 100000,
+            "prefetch_count": 5,
+            "load_all_vertex_attributes": load_all_vertex_attributes,
+            "load_all_edge_attributes": load_all_edge_attributes,
+        }
+
+        if username:
+            db_config_options["username"] = username
+        if password:
+            db_config_options["password"] = password
+        if user_jwt:
+            db_config_options["jwt_token"] = user_jwt
+        if tls_cert:
+            db_config_options["tls_cert"] = tls_cert
+
+        graph_config = {
+            "load_adj_dict": load_adj_dict,
+            "load_coo": load_coo,
+            "is_directed": is_directed,
+            "is_multigraph": is_multigraph,
+            "symmetrize_edges_if_directed": symmetrize_edges_if_directed,
+        }
+
+        vertex_collections = [
+            {"name": v_col_name, "fields": list(entries)}
+            for v_col_name, entries in metagraph["vertexCollections"].items()
+        ]
+
+        edge_collections = [
+            {"name": e_col_name, "fields": list(entries)}
+            for e_col_name, entries in metagraph["edgeCollections"].items()
+        ]
+
+        (
+            node_dict,
+            adj_dict,
+            src_indices,
+            dst_indices,
+            edge_indices,
+            id_to_index_map,
+            edge_values,
+        ) = graph_to_networkx_format(
+            request={
+                "vertex_collections": vertex_collections,
+                "edge_collections": edge_collections,
+                "database_config": db_config_options,
+                "load_config": load_config_options,
+            },
+            graph_config=graph_config,  # TODO Anthony: Move into request
+        )
+
+        return (
+            node_dict,
+            adj_dict,
+            src_indices,
+            dst_indices,
+            edge_indices,
+            id_to_index_map,
+            edge_values,
+        )

phenolrs/networkx/typings.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+Json = dict[str, Any]
+NodeDict = dict[str, Json]
+GraphAdjDict = dict[str, dict[str, Json]]
+DiGraphAdjDict = dict[str, GraphAdjDict]
+MultiGraphAdjDict = dict[str, dict[str, dict[int, Json]]]
+MultiDiGraphAdjDict = dict[str, MultiGraphAdjDict]
+EdgeValuesDict = dict[str, list[int | float]]
+
+SrcIndices = npt.NDArray[np.int64]
+DstIndices = npt.NDArray[np.int64]
+EdgeIndices = npt.NDArray[np.int64]
+ArangoIDtoIndex = dict[str, int]

phenolrs/numpy/__init__.py

@@ -0,0 +1 @@
+from .loader import NumpyLoader  # noqa: F401

phenolrs/numpy/loader.py

@@ -0,0 +1,114 @@
+from typing import Any, Tuple
+
+from phenolrs import PhenolError, graph_to_numpy_format
+
+from .typings import (
+    ArangoCollectionSourceToOutput,
+    ArangoCollectionToArangoKeyToIndex,
+    ArangoCollectionToIndexToArangoKey,
+    ArangoCollectionToNodeFeatures,
+    COOByEdgeType,
+)
+
+
+class NumpyLoader:
+    @staticmethod
+    def load_graph_to_numpy(
+        database: str,
+        metagraph: dict[str, Any],
+        hosts: list[str],
+        user_jwt: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        tls_cert: Any | None = None,
+        parallelism: int | None = None,
+        batch_size: int | None = None,
+    ) -> Tuple[
+        ArangoCollectionToNodeFeatures,
+        COOByEdgeType,
+        ArangoCollectionToArangoKeyToIndex,
+        ArangoCollectionToIndexToArangoKey,
+        ArangoCollectionSourceToOutput,
+    ]:
+        # TODO: replace with pydantic validation
+        db_config_options: dict[str, Any] = {
+            "endpoints": hosts,
+            "database": database,
+        }
+        load_config_options: dict[str, Any] = {
+            "parallelism": parallelism if parallelism is not None else 8,
+            "batch_size": batch_size if batch_size is not None else 100000,
+            "prefetch_count": 5,
+            "load_all_vertex_attributes": False,
+            "load_all_edge_attributes": False,
+        }
+        if username:
+            db_config_options["username"] = username
+        if password:
+            db_config_options["password"] = password
+        if user_jwt:
+            db_config_options["jwt_token"] = user_jwt
+        if tls_cert:
+            db_config_options["tls_cert"] = tls_cert
+
+        if "vertexCollections" not in metagraph:
+            raise PhenolError("vertexCollections not found in metagraph")
+
+        # Address the possibility of having something like this:
+        # "USER": {"x": {"features": None}}
+        # Should be converted to:
+        # "USER": {"x": "features"}
+        entries: dict[str, Any]
+        for v_col_name, entries in metagraph["vertexCollections"].items():
+            for source_name, value in entries.items():
+                if isinstance(value, dict):
+                    if len(value) != 1:
+                        m = f"Only one feature field should be specified per attribute. Found {value}"  # noqa: E501
+                        raise PhenolError(m)
+
+                    value_key = list(value.keys())[0]
+                    if value[value_key] is not None:
+                        m = f"Invalid value for feature {source_name}: {value_key}. Found {value[value_key]}"  # noqa: E501
+                        raise PhenolError(m)
+
+                    metagraph["vertexCollections"][v_col_name][source_name] = value_key
+
+        vertex_collections = [
+            {"name": v_col_name, "fields": list(entries.values())}
+            for v_col_name, entries in metagraph["vertexCollections"].items()
+        ]
+        vertex_cols_source_to_output = {
+            v_col_name: {
+                source_name: output_name for output_name, source_name in entries.items()
+            }
+            for v_col_name, entries in metagraph["vertexCollections"].items()
+        }
+
+        edge_collections = []
+        if "edgeCollections" in metagraph:
+            edge_collections = [
+                {"name": e_col_name, "fields": list(entries.values())}
+                for e_col_name, entries in metagraph["edgeCollections"].items()
+            ]
+
+        (
+            features_by_col,
+            coo_map,
+            col_to_adb_key_to_ind,
+            col_to_ind_to_adb_key,
+        ) = graph_to_numpy_format(
+            {
+                "vertex_collections": vertex_collections,
+                "edge_collections": edge_collections,
+                "database_config": db_config_options,
+                "load_config": load_config_options,
+            }
+        )
+
+        return (
+            features_by_col,
+            coo_map,
+            col_to_adb_key_to_ind,
+            col_to_ind_to_adb_key,
+            vertex_cols_source_to_output,
+        )

phenolrs/numpy/typings.py

@@ -0,0 +1,12 @@
+from typing import Tuple
+
+import numpy as np
+import numpy.typing as npt
+
+EdgeType = Tuple[str, str, str]
+
+ArangoCollectionToNodeFeatures = dict[str, dict[str, npt.NDArray[np.float64]]]
+COOByEdgeType = dict[EdgeType, npt.NDArray[np.float64]]
+ArangoCollectionToArangoKeyToIndex = dict[str, dict[str, int]]
+ArangoCollectionToIndexToArangoKey = dict[str, dict[int, str]]
+ArangoCollectionSourceToOutput = dict[str, dict[str, str]]

phenolrs/phenolrs.pyi

@@ -0,0 +1,84 @@
+import typing
+
+from .networkx.typings import (
+    ArangoIDtoIndex,
+    DiGraphAdjDict,
+    DstIndices,
+    EdgeIndices,
+    EdgeValuesDict,
+    GraphAdjDict,
+    MultiDiGraphAdjDict,
+    MultiGraphAdjDict,
+    NodeDict,
+    SrcIndices,
+)
+from .numpy.typings import (
+    ArangoCollectionToArangoKeyToIndex,
+    ArangoCollectionToIndexToArangoKey,
+    ArangoCollectionToNodeFeatures,
+    COOByEdgeType,
+)
+
+def graph_to_numpy_format(request: dict[str, typing.Any]) -> typing.Tuple[
+    ArangoCollectionToNodeFeatures,
+    COOByEdgeType,
+    ArangoCollectionToArangoKeyToIndex,
+    ArangoCollectionToIndexToArangoKey,
+]: ...
+def graph_to_networkx_format(
+    request: dict[str, typing.Any], graph_config: dict[str, typing.Any]
+) -> typing.Tuple[
+    NodeDict,
+    GraphAdjDict | DiGraphAdjDict | MultiGraphAdjDict | MultiDiGraphAdjDict,
+    SrcIndices,
+    DstIndices,
+    EdgeIndices,
+    ArangoIDtoIndex,
+    EdgeValuesDict,
+]: ...
+
+# AQL-based graph loading types
+class _AqlQueryRequired(typing.TypedDict):
+    """Required fields for AqlQuery."""
+
+    query: str  # The AQL query string
+
+class AqlQuery(_AqlQueryRequired, total=False):
+    """An AQL query with optional bind variables."""
+
+    bindVars: dict[str, typing.Any]  # Optional: bind variables for the query
+
+AttributeSpec = dict[str, str]  # {"attr_name": "type_name"}
+# OR list of {"name": str, "type": str}
+
+AqlDataLoadRequest = typing.TypedDict(
+    "AqlDataLoadRequest",
+    {
+        "database_config": dict[str, typing.Any],
+        "batch_size": int,
+        "vertex_attributes": AttributeSpec | list[dict[str, str]],
+        "edge_attributes": AttributeSpec | list[dict[str, str]],
+        "queries": list[list[AqlQuery]],
+    },
+    total=False,
+)
+
+def graph_aql_to_numpy_format(request: AqlDataLoadRequest) -> typing.Tuple[
+    ArangoCollectionToNodeFeatures,
+    COOByEdgeType,
+    ArangoCollectionToArangoKeyToIndex,
+    ArangoCollectionToIndexToArangoKey,
+]: ...
+def graph_aql_to_networkx_format(
+    request: AqlDataLoadRequest, graph_config: dict[str, typing.Any]
+) -> typing.Tuple[
+    NodeDict,
+    GraphAdjDict | DiGraphAdjDict | MultiGraphAdjDict | MultiDiGraphAdjDict,
+    SrcIndices,
+    DstIndices,
+    EdgeIndices,
+    ArangoIDtoIndex,
+    EdgeValuesDict,
+]: ...
+
+class PhenolError(Exception): ...

phenolrs/pyg/__init__.py

@@ -0,0 +1 @@
+from .loader import PygLoader  # noqa: F401

phenolrs/pyg/loader.py

@@ -0,0 +1,173 @@
+from typing import Any, Tuple
+
+import numpy as np
+
+from phenolrs import PhenolError
+from phenolrs.numpy import NumpyLoader
+
+from .typings import (
+    ArangoCollectionToArangoKeyToIndex,
+    ArangoCollectionToIndexToArangoKey,
+)
+
+try:
+    import torch
+    from torch_geometric.data import Data, HeteroData
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+
+
+class PygLoader:
+    @staticmethod
+    def load_into_pyg_data(
+        database: str,
+        metagraph: dict[str, Any],
+        hosts: list[str],
+        user_jwt: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        tls_cert: Any | None = None,
+        parallelism: int | None = None,
+        batch_size: int | None = None,
+    ) -> Tuple[
+        "Data", ArangoCollectionToArangoKeyToIndex, ArangoCollectionToIndexToArangoKey
+    ]:
+        if not TORCH_AVAILABLE:
+            m = "Missing required dependencies. Install with `pip install phenolrs[torch]`"  # noqa: E501
+            raise ImportError(m)
+
+        if "vertexCollections" not in metagraph:
+            raise PhenolError("vertexCollections not found in metagraph")
+        if "edgeCollections" not in metagraph:
+            raise PhenolError("edgeCollections not found in metagraph")
+
+        if len(metagraph["vertexCollections"]) == 0:
+            raise PhenolError("vertexCollections must map to non-empty dictionary")
+        if len(metagraph["edgeCollections"]) == 0:
+            raise PhenolError("edgeCollections must map to non-empty dictionary")
+
+        if len(metagraph["vertexCollections"]) > 1:
+            m = "More than one vertex collection specified for homogeneous dataset"
+            raise PhenolError(m)
+        if len(metagraph["edgeCollections"]) > 1:
+            m = "More than one edge collection specified for homogeneous dataset"
+            raise PhenolError(m)
+
+        v_col_spec_name = list(metagraph["vertexCollections"].keys())[0]
+        v_col_spec = list(metagraph["vertexCollections"].values())[0]
+
+        (
+            features_by_col,
+            coo_map,
+            col_to_adb_key_to_ind,
+            col_to_ind_to_adb_key,
+            vertex_cols_source_to_output,
+        ) = NumpyLoader.load_graph_to_numpy(
+            database,
+            metagraph,
+            hosts,
+            user_jwt,
+            username,
+            password,
+            tls_cert,
+            parallelism,
+            batch_size,
+        )
+
+        data = Data()
+        # add the features
+        if v_col_spec_name not in features_by_col:
+            raise PhenolError(f"Unable to load data for collection {v_col_spec_name}")
+        for feature in v_col_spec.keys():
+            feature_source_key = v_col_spec[feature]
+            if feature_source_key not in features_by_col[v_col_spec_name]:
+                raise PhenolError(
+                    f"Unable to load features {feature_source_key} for collection {v_col_spec_name}"  # noqa: E501
+                )
+            result = torch.from_numpy(
+                features_by_col[v_col_spec_name][feature_source_key].astype(np.float64)
+            )
+            if result.numel() > 0:
+                data[feature] = result
+
+        # finally add the edges
+        edge_col_name = list(metagraph["edgeCollections"].keys())[0]
+        for e_tup in coo_map.keys():
+            e_name, from_name, to_name = e_tup
+            if e_name == edge_col_name:
+                result = torch.from_numpy(coo_map[e_tup].astype(np.int64))
+                if result.numel() > 0:
+                    data["edge_index"] = result
+
+        return data, col_to_adb_key_to_ind, col_to_ind_to_adb_key
+
+    @staticmethod
+    def load_into_pyg_heterodata(
+        database: str,
+        metagraph: dict[str, Any],
+        hosts: list[str],
+        user_jwt: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        tls_cert: Any | None = None,
+        parallelism: int | None = None,
+        batch_size: int | None = None,
+    ) -> tuple[
+        "HeteroData",
+        ArangoCollectionToArangoKeyToIndex,
+        ArangoCollectionToIndexToArangoKey,
+    ]:
+        if not TORCH_AVAILABLE:
+            m = "Missing required dependencies. Install with `pip install phenolrs[torch]`"  # noqa: E501
+            raise ImportError(m)
+
+        if "vertexCollections" not in metagraph:
+            raise PhenolError("vertexCollections not found in metagraph")
+        if "edgeCollections" not in metagraph:
+            raise PhenolError("edgeCollections not found in metagraph")
+
+        if len(metagraph["vertexCollections"]) == 0:
+            raise PhenolError("vertexCollections must map to non-empty dictionary")
+        if len(metagraph["edgeCollections"]) == 0:
+            raise PhenolError("edgeCollections must map to non-empty dictionary")
+
+        (
+            features_by_col,
+            coo_map,
+            col_to_adb_key_to_ind,
+            col_to_ind_to_adb_key,
+            vertex_cols_source_to_output,
+        ) = NumpyLoader.load_graph_to_numpy(
+            database,
+            metagraph,
+            hosts,
+            user_jwt,
+            username,
+            password,
+            tls_cert,
+            parallelism,
+            batch_size,
+        )
+        data = HeteroData()
+        for col in features_by_col.keys():
+            col_mapping = vertex_cols_source_to_output[col]
+            for feature in features_by_col[col].keys():
+                if feature == "@collection_name":
+                    continue
+
+                target_name = col_mapping[feature]
+                result = torch.from_numpy(
+                    features_by_col[col][feature].astype(np.float64)
+                )
+                if result.numel() > 0:
+                    data[col][target_name] = result
+
+        for edge_col in coo_map.keys():
+            edge_col_name, from_name, to_name = edge_col
+            result = torch.from_numpy(coo_map[edge_col].astype(np.int64))
+            if result.numel() > 0:
+                data[(from_name, edge_col_name, to_name)].edge_index = result
+
+        return data, col_to_adb_key_to_ind, col_to_ind_to_adb_key

phenolrs/pyg/typings.py

@@ -0,0 +1,2 @@
+ArangoCollectionToArangoKeyToIndex = dict[str, dict[str, int]]
+ArangoCollectionToIndexToArangoKey = dict[str, dict[int, str]]

phenolrs-0.5.12.dist-info/METADATA

@@ -0,0 +1,24 @@
+Metadata-Version: 2.4
+Name: phenolrs
+Version: 0.5.12
+Classifier: Programming Language :: Rust
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+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: Typing :: Typed
+Requires-Dist: numpy
+Requires-Dist: python-arango>=8.2.6
+Requires-Dist: pytest ; extra == 'tests'
+Requires-Dist: arango-datasets ; extra == 'tests'
+Requires-Dist: adbnx-adapter ; extra == 'tests'
+Requires-Dist: torch ; extra == 'torch'
+Requires-Dist: torch-geometric ; extra == 'torch'
+Requires-Dist: version ; extra == 'dynamic'
+Provides-Extra: tests
+Provides-Extra: torch
+Provides-Extra: dynamic
+License-File: LICENSE
+Requires-Python: >=3.10