rxgraph 0.13.0__cp315-cp315-win_amd64.whl

rxgraph/__init__.py

@@ -0,0 +1,4 @@
+from . import _rxgraph
+from .plugin import export_api
+
+export_api(globals(), _rxgraph)

rxgraph/__init__.pyi

@@ -0,0 +1,310 @@
+from collections.abc import Hashable, Iterable, Mapping
+from dataclasses import dataclass
+from os import PathLike
+from typing import Any, Literal, Self, TypeVar
+
+from polars import DataFrame, Expr, LazyFrame, col as col, lit as lit
+
+Node = TypeVar("Node", bound=Hashable)
+GraphId = int | str
+
+def rayon_thread_count() -> int:
+    """Return the Rayon worker thread count used by rxgraph."""
+    ...
+
+class Graph:
+    """Arrow-backed directed graph."""
+
+    def __init__(
+        self,
+        nodes: DataFrame | list[tuple[str, DataFrame]],
+        edges: DataFrame | list[tuple[str, DataFrame]],
+    ) -> None:
+        """Build a graph from node and edge tables.
+
+        Nodes require ``id`` and edges require ``id``, ``src``, and ``dest``.
+        All ID columns must be either UInt64 or string. Optional ``type``
+        columns must be string when present.
+        """
+        ...
+
+    @classmethod
+    def from_edges(
+        cls,
+        edges: Iterable[tuple[Node, Node] | tuple[Node, Node, Mapping[str, Any]]],
+        *,
+        nodes: Iterable[Node | tuple[Node, Mapping[str, Any]]] | None = None,
+    ) -> Self:
+        """Build a directed graph from Python node labels and edge tuples."""
+        ...
+
+    @classmethod
+    def from_lazy(cls, nodes: LazyFrame, edges: LazyFrame) -> Self:
+        """Build a graph from Polars ``LazyFrame``s, deferring payload columns.
+
+        Only identity/topology columns (``id`` for nodes; ``id``/``src``/``dest``
+        for edges, plus optional ``type``) are collected at construction. Payload
+        columns are pulled lazily at :meth:`search` time, projected to only the
+        columns the kernel references, bounding resident payload memory.
+        """
+        ...
+
+    @classmethod
+    def from_parquet(
+        cls,
+        nodes: str | PathLike[str],
+        edges: str | PathLike[str],
+        *,
+        payloads: Literal["eager", "lazy"] = "eager",
+    ) -> Self:
+        """Build a file-backed graph from Parquet node and edge tables."""
+        ...
+
+    @property
+    def node_count(self) -> int: ...
+    @property
+    def edge_count(self) -> int: ...
+    def node_id(self, label: Hashable) -> GraphId:
+        """Return the graph ID used by the engine for a node label."""
+        ...
+    def search(
+        self,
+        *,
+        start_nodes: Iterable[Hashable],
+        visit: Expr | str | None = None,
+        next_state: Mapping[str, Expr | str] | None = None,
+        stop: Expr | str | None = None,
+        initial_state: Mapping[str, Any] | None = None,
+        kernel: str | None = None,
+        params: Mapping[str, Any] | None = None,
+        columns: Iterable[str] | None = None,
+        max_depth: int | None = None,
+        max_paths: int | None = None,
+        strategy: Literal["dfs", "bfs"] = "dfs",
+        parallel: bool | Literal["auto", "off", "on"] = True,
+        intermediate_states: bool = False,
+        progress: bool = False,
+        max_visits_per_node: int = 1,
+    ) -> SearchResult:
+        """Run a stateful traversal.
+
+        Two mutually exclusive modes:
+
+        * DSL: ``visit``/``next_state``/``stop``/``initial_state`` Polars expressions.
+        * Named native kernel: ``kernel`` (registered Rust kernel name) and ``params``.
+          The DSL args must not be set in this mode.
+
+        ``params`` node-label values under the conventional keys
+        ``target``/``source``/``start``/``node`` are translated to engine IDs when
+        the value is a known node label; other values pass through unchanged.
+
+        ``columns`` lists payload columns to load for ``from_lazy`` before a
+        legacy named kernel runs. File-backed ``from_parquet`` graphs use typed
+        native kernels and do not need ``columns``.
+        """
+        ...
+    def search_batches(
+        self,
+        *,
+        start_nodes: Iterable[Hashable],
+        kernel: str,
+        params: Mapping[str, Any] | None = None,
+        columns: Iterable[str] | None = None,
+        max_depth: int | None = None,
+        max_paths: int | None = None,
+        strategy: Literal["dfs", "bfs"] = "dfs",
+        parallel: bool | Literal["auto", "off", "on"] = True,
+        intermediate_states: bool = False,
+        progress: bool = False,
+        max_visits_per_node: int = 1,
+        batch_size: int = 1024,
+    ) -> SearchStream: ...
+    def search_paths(
+        self,
+        *,
+        start_nodes: Iterable[Hashable],
+        visit: Expr | str | None = None,
+        next_state: Mapping[str, Expr | str] | None = None,
+        stop: Expr | str | None = None,
+        initial_state: Mapping[str, Any] | None = None,
+        kernel: str | None = None,
+        params: Mapping[str, Any] | None = None,
+        columns: Iterable[str] | None = None,
+        max_depth: int | None = None,
+        max_paths: int | None = None,
+        strategy: Literal["dfs", "bfs"] = "dfs",
+        parallel: bool | Literal["auto", "off", "on"] = True,
+        intermediate_states: bool = False,
+        progress: bool = False,
+        max_visits_per_node: int = 1,
+    ) -> SearchResult: ...
+    def search_first(
+        self,
+        *,
+        start_nodes: Iterable[Hashable],
+        visit: Expr | str | None = None,
+        next_state: Mapping[str, Expr | str] | None = None,
+        stop: Expr | str | None = None,
+        initial_state: Mapping[str, Any] | None = None,
+        kernel: str | None = None,
+        params: Mapping[str, Any] | None = None,
+        columns: Iterable[str] | None = None,
+        max_depth: int | None = None,
+        intermediate_states: bool = False,
+        progress: bool = False,
+        max_visits_per_node: int = 1,
+    ) -> SearchPath | None: ...
+    def bfs(self, start: Hashable, max_depth: int | None = None) -> list[Any]:
+        """Return nodes reachable from ``start`` in breadth-first order."""
+        ...
+    def dfs(self, start: Hashable, max_depth: int | None = None) -> list[Any]:
+        """Return nodes reachable from ``start`` in depth-first pre-order."""
+        ...
+    def reachable_nodes(self, start: Hashable) -> list[Any]:
+        """Return all nodes reachable from ``start``."""
+        ...
+    def ancestors_many(self, targets: Iterable[Hashable]) -> list[Any]:
+        """Return nodes that can reach any target, excluding targets, in BFS order."""
+        ...
+    def descendants_many(self, sources: Iterable[Hashable]) -> list[Any]:
+        """Return nodes reachable from any source, excluding sources, in BFS order."""
+        ...
+    def shortest_path(self, source: Hashable, target: Hashable) -> list[Any] | None:
+        """Return an unweighted directed shortest path, if one exists."""
+        ...
+    def out_degrees(self) -> list[int]:
+        """Return out-degree for each node in node insertion order."""
+        ...
+    def in_degrees(self) -> list[int]:
+        """Return in-degree for each node in node insertion order."""
+        ...
+    def degrees(self) -> list[int]:
+        """Return total directed degree for each node in node insertion order."""
+        ...
+    def weakly_connected_components(self) -> list[list[Any]]:
+        """Return weakly connected components, ignoring edge direction."""
+        ...
+
+class DiGraph(Graph):
+    """Arrow-backed bidirectional graph."""
+
+    def __init__(
+        self,
+        nodes: DataFrame | list[tuple[str, DataFrame]],
+        edges: DataFrame | list[tuple[str, DataFrame]],
+    ) -> None:
+        """Build a bidirectional graph from node and edge tables.
+
+        Each input edge is available in both directions. Reverse rows receive
+        synthetic engine IDs, but traversal path edges are mapped back to the
+        original input edge IDs.
+        """
+        ...
+
+    @classmethod
+    def from_edges(
+        cls,
+        edges: Iterable[tuple[Node, Node] | tuple[Node, Node, Mapping[str, Any]]],
+        *,
+        nodes: Iterable[Node | tuple[Node, Mapping[str, Any]]] | None = None,
+    ) -> Self:
+        """Build a bidirectional graph from Python node labels and edge tuples."""
+        ...
+
+class Kernel:
+    """Traversal kernel built from Polars expressions."""
+
+    visit: Expr | str
+    next_state: dict[str, Expr | str]
+    stop: Expr | str
+    initial_state: dict[str, Any]
+
+    def __init__(
+        self,
+        visit: Expr | str | None = None,
+        next_state: Mapping[str, Expr | str] | None = None,
+        stop: Expr | str | None = None,
+        initial_state: Mapping[str, Any] | None = None,
+    ) -> None:
+        """Create a kernel.
+
+        ``visit`` decides whether an edge is accepted, ``next_state`` computes
+        state updates for accepted edges, and ``stop`` decides whether an
+        accepted edge materializes a path. Defaults accept every edge, leave
+        state unchanged, never stop, and start with empty state.
+        """
+        ...
+
+class Traversal:
+    """Low-level traversal configuration for the native engine."""
+
+    def __init__(
+        self,
+        kernel: Kernel,
+        start_nodes: list[Hashable],
+        max_depth: int,
+        max_paths: int,
+        strategy: Literal["dfs", "bfs"] = "dfs",
+        parallel: bool | Literal["auto", "off", "on"] = True,
+        intermediate_states: bool = False,
+        progress: bool = False,
+    ) -> None:
+        """Create a traversal.
+
+        ``strategy`` is ``"dfs"`` or ``"bfs"``. ``parallel`` enables or disables
+        rxgraph parallel traversal. ``intermediate_states`` stores per-node
+        state history on returned paths.
+        """
+        ...
+
+class SearchStats:
+    """Counters collected while searching."""
+
+    @property
+    def start_nodes(self) -> int: ...
+    @property
+    def path_entries(self) -> int: ...
+    @property
+    def evaluated_edges(self) -> int: ...
+    @property
+    def parallel_edges(self) -> int: ...
+    @property
+    def accepted_edges(self) -> int: ...
+    @property
+    def stopped_paths(self) -> int: ...
+    @property
+    def rejected_edges(self) -> int: ...
+    @property
+    def skipped_revisits(self) -> int: ...
+    @property
+    def max_depth(self) -> int: ...
+    @property
+    def materialized_node_payloads(self) -> int: ...
+    @property
+    def materialized_edge_payloads(self) -> int: ...
+
+@dataclass(slots=True)
+class SearchPath:
+    """One stopped path returned by a traversal."""
+
+    nodes: list[Any]
+    edges: list[Any]
+    state: dict[str, Any]
+    intermediate_states: list[dict[str, Any]] | None = None
+
+@dataclass(slots=True)
+class SearchResult:
+    """Paths and stats returned by :meth:`Graph.search`."""
+
+    paths: list[SearchPath]
+    stats: SearchStats
+
+class SearchStream:
+    """Native path batches; use close() or a context manager for early termination."""
+    def __iter__(self) -> SearchStream: ...
+    def __next__(self) -> list[SearchPath]: ...
+    @property
+    def stats(self) -> SearchStats: ...
+    def close(self) -> None: ...
+    def __enter__(self) -> SearchStream: ...
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: ...

rxgraph/_graph_tables.py

@@ -0,0 +1,219 @@
+from collections.abc import Hashable, Iterable, Mapping
+from dataclasses import dataclass
+from typing import Any, cast
+
+import polars as pl
+
+ID_COL = "id"
+SRC_COL = "src"
+DEST_COL = "dest"
+NODE_RESERVED_ATTRS = frozenset({ID_COL})
+EDGE_RESERVED_ATTRS = frozenset({ID_COL, SRC_COL, DEST_COL})
+REVERSE_EDGE_SUFFIX = "__rev"
+UINT64_MAX = (1 << 64) - 1
+
+NodeInput = Hashable | tuple[Hashable, Mapping[str, Any]]
+EdgeInput = tuple[Hashable, Hashable] | tuple[Hashable, Hashable, Mapping[str, Any]]
+
+
+@dataclass(frozen=True)
+class GraphTables:
+    nodes: pl.DataFrame
+    edges: pl.DataFrame
+    label_to_id: dict[Hashable, int]
+    id_to_label: list[Hashable]
+    edge_id_to_label: dict[Hashable, Hashable] | None
+
+
+def normalize_table(value: Any) -> Any:
+    if isinstance(value, list):
+        if len(value) != 1:
+            raise ValueError(
+                "rxgraph expects one node DataFrame and one edge DataFrame"
+            )
+        return value[0][1]
+    return value
+
+
+def build_labeled_tables(
+    edges: Iterable[EdgeInput],
+    nodes: Iterable[NodeInput] | None,
+    *,
+    bidirectional: bool = False,
+) -> GraphTables:
+    builder = _LabeledTableBuilder(track_edge_labels=bidirectional)
+
+    if nodes is not None:
+        for node in nodes:
+            label, attrs = _parse_node(node)
+            builder.add_node(label, attrs)
+
+    for edge_id, edge in enumerate(edges):
+        src, dest, attrs = _parse_edge(edge)
+        builder.add_edge(src, dest, attrs, edge_id)
+        if bidirectional:
+            builder.add_edge(dest, src, attrs, edge_id)
+
+    return builder.build()
+
+
+def build_bidirectional_edges(
+    edges: Any,
+) -> tuple[pl.DataFrame, dict[Hashable, Hashable]]:
+    edge_table = normalize_table(edges)
+    if not isinstance(edge_table, pl.DataFrame):
+        edge_table = pl.DataFrame(edge_table)
+
+    edge_ids = edge_table[ID_COL].to_list()
+    reverse_ids = _reverse_edge_ids(edge_ids, edge_table.schema[ID_COL])
+    reverse_table = edge_table.with_columns(
+        pl.Series(ID_COL, reverse_ids, dtype=edge_table.schema[ID_COL]),
+        pl.col(DEST_COL).alias(SRC_COL),
+        pl.col(SRC_COL).alias(DEST_COL),
+    )
+    edge_id_to_label = {edge_id: edge_id for edge_id in edge_ids}
+    edge_id_to_label.update(
+        {
+            reverse_id: edge_id
+            for reverse_id, edge_id in zip(reverse_ids, edge_ids, strict=True)
+        }
+    )
+    return pl.concat([edge_table, reverse_table], how="vertical"), edge_id_to_label
+
+
+class _LabeledTableBuilder:
+    label_to_id: dict[Hashable, int]
+    id_to_label: list[Hashable]
+    node_attrs: list[dict[str, Any]]
+    edge_srcs: list[int]
+    edge_dests: list[int]
+    edge_attrs: list[dict[str, Any]]
+    edge_id_to_label: dict[Hashable, Hashable] | None
+
+    def __init__(self, *, track_edge_labels: bool) -> None:
+        self.label_to_id = {}
+        self.id_to_label = []
+        self.node_attrs = []
+        self.edge_srcs = []
+        self.edge_dests = []
+        self.edge_attrs = []
+        self.edge_id_to_label = {} if track_edge_labels else None
+
+    def add_node(self, label: Hashable, attrs: Mapping[str, Any] | None = None) -> int:
+        if label in self.label_to_id:
+            if attrs:
+                self.node_attrs[self.label_to_id[label]].update(_attrs(attrs, "node"))
+            return self.label_to_id[label]
+
+        node_id = len(self.id_to_label)
+        self.label_to_id[label] = node_id
+        self.id_to_label.append(label)
+        self.node_attrs.append(_attrs(attrs or {}, "node"))
+        return node_id
+
+    def add_edge(
+        self,
+        src: Hashable,
+        dest: Hashable,
+        attrs: Mapping[str, Any] | None,
+        label: Hashable,
+    ) -> None:
+        edge_id = len(self.edge_srcs)
+        self.edge_srcs.append(self.add_node(src))
+        self.edge_dests.append(self.add_node(dest))
+        self.edge_attrs.append(_attrs(attrs or {}, "edge"))
+        if self.edge_id_to_label is not None:
+            self.edge_id_to_label[edge_id] = label
+
+    def build(self) -> GraphTables:
+        node_data = _rows_to_columns(self.node_attrs)
+        node_data[ID_COL] = list(range(len(self.id_to_label)))
+        edge_data = _rows_to_columns(self.edge_attrs)
+        edge_data[ID_COL] = list(range(len(self.edge_srcs)))
+        edge_data[SRC_COL] = self.edge_srcs
+        edge_data[DEST_COL] = self.edge_dests
+
+        return GraphTables(
+            nodes=pl.DataFrame(node_data, schema_overrides={ID_COL: pl.UInt64}),
+            edges=pl.DataFrame(
+                edge_data,
+                schema_overrides={
+                    ID_COL: pl.UInt64,
+                    SRC_COL: pl.UInt64,
+                    DEST_COL: pl.UInt64,
+                },
+            ),
+            label_to_id=self.label_to_id,
+            id_to_label=self.id_to_label,
+            edge_id_to_label=self.edge_id_to_label,
+        )
+
+
+def _parse_node(node: NodeInput) -> tuple[Hashable, Mapping[str, Any] | None]:
+    if isinstance(node, tuple) and len(node) == 2 and isinstance(node[1], Mapping):
+        return node[0], cast(Mapping[str, Any], node[1])
+    return node, None
+
+
+def _parse_edge(edge: EdgeInput) -> tuple[Hashable, Hashable, Mapping[str, Any] | None]:
+    if len(edge) == 2:
+        return edge[0], edge[1], None
+    attrs = cast(tuple[Hashable, Hashable, Any], edge)[2]
+    if isinstance(attrs, Mapping):
+        return edge[0], edge[1], cast(Mapping[str, Any], attrs)
+    raise ValueError("edges must be (src, dest) or (src, dest, attrs) tuples")
+
+
+def _attrs(attrs: Mapping[str, Any], kind: str) -> dict[str, Any]:
+    reserved = NODE_RESERVED_ATTRS if kind == "node" else EDGE_RESERVED_ATTRS
+    overlap = reserved.intersection(attrs)
+    if overlap:
+        names = ", ".join(sorted(overlap))
+        raise ValueError(f"{kind} attributes cannot use reserved keys: {names}")
+    return dict(attrs)
+
+
+def _rows_to_columns(rows: list[dict[str, Any]]) -> dict[str, list[Any]]:
+    keys = sorted(
+        {key for row in rows for key, value in row.items() if value is not None}
+    )
+    return {key: [row.get(key) for row in rows] for key in keys}
+
+
+def _reverse_edge_ids(edge_ids: list[Any], dtype: pl.DataType) -> list[Any]:
+    used = set(edge_ids)
+    if dtype.is_integer():
+        return _reverse_integer_edge_ids(edge_ids, used)
+    if dtype == pl.String:
+        return _reverse_string_edge_ids(edge_ids, used)
+    raise ValueError("DiGraph edge ids must be integers or strings")
+
+
+def _reverse_integer_edge_ids(edge_ids: list[Any], used: set[Any]) -> list[int]:
+    max_id = max((edge_id for edge_id in edge_ids if edge_id is not None), default=-1)
+    next_id = max_id + 1
+    if next_id + len(edge_ids) - 1 <= UINT64_MAX:
+        return list(range(next_id, next_id + len(edge_ids)))
+
+    reverse_ids: list[int] = []
+    candidate = 0
+    for _ in edge_ids:
+        while candidate in used:
+            candidate += 1
+        reverse_ids.append(candidate)
+        used.add(candidate)
+    return reverse_ids
+
+
+def _reverse_string_edge_ids(edge_ids: list[Any], used: set[Any]) -> list[str]:
+    reverse_ids: list[str] = []
+    for edge_id in edge_ids:
+        base = f"{edge_id}{REVERSE_EDGE_SUFFIX}"
+        candidate = base
+        suffix = 2
+        while candidate in used:
+            candidate = f"{base}_{suffix}"
+            suffix += 1
+        reverse_ids.append(candidate)
+        used.add(candidate)
+    return reverse_ids