diff --git a/src/omop_graph/graph/kg.py b/src/omop_graph/graph/kg.py index 0072510..a0b0879 100644 --- a/src/omop_graph/graph/kg.py +++ b/src/omop_graph/graph/kg.py @@ -605,12 +605,14 @@ def iter_edges( active_only: bool = True, on: Optional[date] = None, within_domain: bool = False, + stream: bool = False, ) -> Generator[EdgeView, None, None]: """Yield edges of *concept_ids* on *session*, classified from the relationship mapping. A *predicate_kinds* filter is applied in SQL as the relationship IDs of those kinds. Edges whose relationship has no classification are - skipped. + skipped. Set ``stream=True`` to fetch rows in batches instead of + buffering the full result set. """ if predicate_kinds: kind_ids = frozenset( @@ -630,15 +632,23 @@ def iter_edges( on=on, within_domain=within_domain, ) - for row in session.execute(stmt): - mapping = self._relationship_mapping.get(row.predicate_id) - if mapping is None: - continue - yield EdgeView( - **row._mapping, - predicate_kind=mapping.predicate_kind, - predicate_subkind=mapping.predicate_subkind, - ) + result = ( + session.execute(stmt, execution_options={"yield_per": 1000}) + if stream + else session.execute(stmt) + ) + try: + for row in result: + mapping = self._relationship_mapping.get(row.predicate_id) + if mapping is None: + continue + yield EdgeView( + **row._mapping, + predicate_kind=mapping.predicate_kind, + predicate_subkind=mapping.predicate_subkind, + ) + finally: + result.close() def specificity(self, concept_id: int) -> float: """ diff --git a/src/omop_graph/graph/traverse.py b/src/omop_graph/graph/traverse.py index efa995b..c8baf87 100644 --- a/src/omop_graph/graph/traverse.py +++ b/src/omop_graph/graph/traverse.py @@ -12,6 +12,7 @@ from __future__ import annotations from collections import deque +from contextlib import closing from dataclasses import dataclass from datetime import date from typing import TYPE_CHECKING, Iterable, List, Optional, Set, Tuple @@ -111,6 +112,48 @@ def summary(self, kg: "KnowledgeGraph", max_steps: int = 10) -> str: return "\n".join(lines) +@dataclass +class _TraversalState: + queue: deque[Tuple[int, int]] + discovered: set[int] + max_nodes: Optional[int] + predicate_kinds: Optional[Set[PredicateKind]] + on: Optional[date] + + +def _expand_node( + kg: "KnowledgeGraph", + node: int, + depth: int, + state: _TraversalState, +) -> Tuple[List[EdgeView], bool]: + """Collect one node's edges, stopping when a new target fills the budget.""" + expanded: List[EdgeView] = [] + with kg.session_factory() as session: + edge_iter = kg.iter_edges( + session=session, + concept_ids=node, + direction="out", + predicate_kinds=frozenset(state.predicate_kinds) + if state.predicate_kinds is not None + else None, + active_only=True, + on=state.on, + stream=True, + ) + with closing(edge_iter): + for edge in edge_iter: + if edge.object_id not in state.discovered: + state.discovered.add(edge.object_id) + state.queue.append((edge.object_id, depth + 1)) + + expanded.append(edge) + if state.max_nodes and len(state.discovered) >= state.max_nodes: + return expanded, True + + return expanded, False + + def traverse( kg: "KnowledgeGraph", seeds: Iterable[int], @@ -136,7 +179,7 @@ def traverse( on : date, optional Filter for edges active on this date. max_nodes : int, optional - Stop after visiting this many unique nodes. + Stop after discovering this many unique nodes. trace : bool If True, return a GraphTrace object. @@ -144,72 +187,58 @@ def traverse( ------- tuple[Subgraph, GraphTrace | None] The resulting subgraph and optionally the execution trace. + + When ``max_nodes`` is reached, which nodes fill the budget is unspecified + (it depends on database row order). Traversal stops as soon as the node + budget is discovered, so queued nodes inside the budget are not expanded; + edges those later expansions could have found are absent from the result. """ # Deduplicate seeds while preserving order # (Python 3.7+ dict guarantees insertion order) unique_seeds = tuple(dict.fromkeys(seeds)) - visited = set() + # Seeds precede discovered neighbours in BFS order. Keep only the seeds + # that fit the budget, matching the order in which the old visited set + # reached its cap. + budgeted_seeds = unique_seeds[: max(max_nodes, 1)] if max_nodes else unique_seeds + discovered = set(budgeted_seeds) edges_out: List[EdgeView] = [] trace_steps: List[TraceStep] = [] # Queue stores (concept_id, depth) - q = deque((s, 0) for s in unique_seeds) + q = deque((s, 0) for s in budgeted_seeds) + state = _TraversalState(q, discovered, max_nodes, predicate_kinds, on) terminated = None while q: node, depth = q.popleft() - if node in visited: - continue - - visited.add(node) - - if max_nodes and len(visited) >= max_nodes: + if state.max_nodes and len(state.discovered) >= state.max_nodes: terminated = "max_nodes" break if depth >= max_depth: continue - expanded: List[EdgeView] = [] - - # Iterate over outgoing edges - with kg.session_factory() as session: - for e in kg.iter_edges( - session=session, - concept_ids=node, - direction="out", - predicate_kinds=frozenset(predicate_kinds) - if predicate_kinds is not None - else None, - active_only=True, - on=on, - ): - # Only add edges to the output; we decide to traverse in the next block - expanded.append(e) - edges_out.append(e) - - nxt = e.object_id - # Optimization: Don't add to queue if already visited - if nxt not in visited: - q.append((nxt, depth + 1)) + # Iterate over outgoing edges until one more node would exceed the budget. + expanded, budget_reached = _expand_node(kg, node, depth, state) + edges_out.extend(expanded) if trace: trace_steps.append( TraceStep(depth=depth, node=node, expanded_edges=tuple(expanded)) ) + if budget_reached: + terminated = "max_nodes" + break - # Deduplicate edges and drop any that target an unvisited node. - # The latter can happen when max_nodes terminates the loop while neighbour - # nodes are still queued but never expanded, which would break the invariant - # that every edge in the Subgraph is closed over its node set. + # Deduplicate edges and drop any that target a node outside the budget. dedup = { (e.subject_id, e.predicate_id, e.object_id): e for e in edges_out - if e.object_id in visited + if e.subject_id in state.discovered and e.object_id in state.discovered } - sg = Subgraph(nodes=frozenset(visited), edges=tuple(dedup.values())) + sg = Subgraph(nodes=frozenset(state.discovered), edges=tuple(dedup.values())) graph_trace = ( GraphTrace(seeds=unique_seeds, steps=trace_steps, terminated_reason=terminated) diff --git a/tests/test_traverse.py b/tests/test_traverse.py new file mode 100644 index 0000000..6117cd4 --- /dev/null +++ b/tests/test_traverse.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +from contextlib import nullcontext +from datetime import date +from types import SimpleNamespace + +from omop_graph.extensions.omop_alchemy import ( + PredicateKind, + RelationshipMappingElement, +) +from omop_graph.graph.edges import EdgeView +from omop_graph.graph.kg import KnowledgeGraph +from omop_graph.graph.traverse import traverse + + +class FakeKnowledgeGraph: + def __init__(self, adjacency: dict[int, tuple[int, ...]]) -> None: + self.adjacency = adjacency + self.materialized_edges = 0 + self.closed_edge_iterators = 0 + self.stream_flags: list[bool] = [] + + @staticmethod + def session_factory(): + return nullcontext(None) + + def iter_edges(self, *, concept_ids: int, stream: bool = False, **_kwargs): + self.stream_flags.append(stream) + try: + for object_id in self.adjacency.get(concept_ids, ()): + self.materialized_edges += 1 + yield EdgeView( + subject_id=concept_ids, + predicate_id="is a", + object_id=object_id, + valid_start_date=date(2000, 1, 1), + valid_end_date=date(2099, 12, 31), + invalid_reason=None, + predicate_kind=PredicateKind.HIERARCHY, + predicate_subkind="test", + ) + finally: + self.closed_edge_iterators += 1 + + +def _edge_keys(subgraph) -> set[tuple[int, str, int]]: + return { + (edge.subject_id, edge.predicate_id, edge.object_id) + for edge in subgraph.edges + } + + +def _distances(adjacency: dict[int, tuple[int, ...]], seed: int) -> dict[int, int]: + distances = {seed: 0} + queue = [seed] + for node in queue: + for target in adjacency.get(node, ()): + if target not in distances: + distances[target] = distances[node] + 1 + queue.append(target) + return distances + + +def test_uncapped_traversal_keeps_existing_nodes_and_edges() -> None: + kg = FakeKnowledgeGraph({1: (2, 3), 2: (4,), 3: (5,), 4: (6,), 5: ()}) + + subgraph, trace = traverse( + kg, + seeds=(1,), + predicate_kinds={PredicateKind.HIERARCHY}, + max_depth=2, + on=None, + max_nodes=20, + trace=True, + ) + + assert subgraph.nodes == frozenset({1, 2, 3, 4, 5}) + assert _edge_keys(subgraph) == { + (1, "is a", 2), + (1, "is a", 3), + (2, "is a", 4), + (3, "is a", 5), + } + assert trace is not None + assert trace.terminated_reason is None + + +def test_capped_traversal_respects_node_depth_and_edge_closure() -> None: + adjacency = {1: (2, 3), 2: (4, 5), 3: (6,), 4: (7,), 5: (), 6: ()} + kg = FakeKnowledgeGraph(adjacency) + + subgraph, trace = traverse( + kg, + seeds=(1,), + predicate_kinds={PredicateKind.HIERARCHY}, + max_depth=2, + on=None, + max_nodes=4, + trace=True, + ) + + assert len(subgraph.nodes) == 4 + assert 1 in subgraph.nodes + distances = _distances(adjacency, 1) + assert all(distances[node] <= 2 for node in subgraph.nodes) + assert all( + edge.subject_id in subgraph.nodes and edge.object_id in subgraph.nodes + for edge in subgraph.edges + ) + assert trace is not None + assert trace.terminated_reason == "max_nodes" + + +def test_fanout_traversal_stops_and_closes_edge_iterator_at_node_budget() -> None: + fanout = 5_000 + kg = FakeKnowledgeGraph({1: tuple(range(2, fanout + 2))}) + + subgraph, _ = traverse( + kg, + seeds=(1,), + predicate_kinds={PredicateKind.HIERARCHY}, + max_depth=2, + on=None, + max_nodes=32, + trace=False, + ) + + assert len(subgraph.nodes) == 32 + assert kg.materialized_edges <= 31 + assert kg.closed_edge_iterators == 1 + assert kg.stream_flags == [True] + + +def test_iter_edges_streaming_is_opt_in() -> None: + class RecordingResult: + def __init__(self, rows=()) -> None: + self.rows = rows + self.closed = False + + def __iter__(self): + return iter(self.rows) + + def close(self) -> None: + self.closed = True + + class RecordingSession: + def __init__(self, results) -> None: + self.execution_options: list[object] = [] + self.results = iter(results) + + def execute(self, _statement, **kwargs): + self.execution_options.append(kwargs.get("execution_options", {})) + return next(self.results) + + kg = KnowledgeGraph.__new__(KnowledgeGraph) + kg._relationship_mapping = { + "is a": RelationshipMappingElement( + "is a", PredicateKind.HIERARCHY, "test" + ) + } + row = SimpleNamespace( + predicate_id="is a", + _mapping={ + "subject_id": 1, + "predicate_id": "is a", + "object_id": 2, + "valid_start_date": date(2000, 1, 1), + "valid_end_date": date(2099, 12, 31), + "invalid_reason": None, + }, + ) + default_result = RecordingResult() + streaming_result = RecordingResult([row, row]) + session = RecordingSession([default_result, streaming_result]) + + list(kg.iter_edges(session=session, concept_ids=1, direction="out")) + edge_iter = kg.iter_edges( + session=session, concept_ids=1, direction="out", stream=True + ) + next(edge_iter) + edge_iter.close() + + assert session.execution_options == [{}, {"yield_per": 1000}] + assert default_result.closed + assert streaming_result.closed