Stop traverse early when max_nodes is reached

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:
         """

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,80 +179,66 @@ 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.
 
     Returns
     -------
     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)

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