diff --git a/src/tracksdata/graph/_base_graph.py b/src/tracksdata/graph/_base_graph.py index 685550c7..ee6aa49a 100644 --- a/src/tracksdata/graph/_base_graph.py +++ b/src/tracksdata/graph/_base_graph.py @@ -4,6 +4,7 @@ from collections.abc import Sequence from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload +from weakref import WeakSet import geff import numpy as np @@ -117,6 +118,14 @@ class BaseGraph(abc.ABC): def __init__(self) -> None: self._cache = {} + # Views derived from this graph, to be kept up to date when it changes. + # Views add themselves on construction (see GraphView.__init__). + # + # Held weakly: once nothing else references a view, nobody can observe + # whether it is current, so maintaining it would be pure overhead. + # Dropping the last reference to a view is all that is needed to stop + # maintaining it. + self._views: WeakSet[BaseGraph] = WeakSet() def supports_custom_indices(self) -> bool: """ @@ -124,6 +133,124 @@ def supports_custom_indices(self) -> bool: """ return False + def __getstate__(self) -> dict[str, Any]: + """ + Drop the view registry when serializing. + + Views are runtime relationships between live objects, not data, so an + unpickled graph starts with none. A ``WeakSet`` is also not picklable, + as it holds an internal callback closure. + """ + state = self.__dict__.copy() + state.pop("_views", None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__.update(state) + # Excluded by __getstate__, so restore it empty: an unpickled graph has + # no live views. + self._views = WeakSet() + + def _views_need_node_attrs(self) -> tuple[bool, bool]: + """ + Whether any registered view needs before/after snapshots of a node update. + + Building those snapshots is the dominant cost of updating a graph that has + views, so it must not be paid for views that cannot use them: a view that + shares this graph's attribute dicts (a rustworkx root) and has no listener + needs nothing at all. + + Returns + ------- + tuple[bool, bool] + ``(needs_old, needs_new)``. Old values are only ever used to emit a + view's ``node_updated``; new values are additionally needed by a view + that keeps its own copy of the attributes and has to be written through. + """ + needs_old = needs_new = False + for view in self._views: + view_old, view_new = view._needs_root_node_attrs() + needs_old |= view_old + needs_new |= view_new + if needs_old and needs_new: + break + return needs_old, needs_new + + def _maintain_views_node_attrs( + self, + node_ids: Sequence[int], + old_attrs_by_id: dict[int, dict[str, Any]] | None, + new_attrs_by_id: dict[int, dict[str, Any]] | None, + changed_keys: set[str], + ) -> None: + """ + Bring every registered view up to date after a node attribute update. + + Called by concrete ``update_node_attrs`` implementations after the root + graph has been updated and its own signal emitted. Each view absorbs + the change for the nodes it contains and then emits its own signals. + + This is deliberately *not* routed through the signal system: keeping a + view consistent with its root is an invariant, so it must happen whether + or not anything is listening to the view. + + `node_ids` and the attribute dicts are keyed by this graph's own node + ids, which is what views map from. + """ + for view in self._views: + view._apply_root_node_attrs( + node_ids=node_ids, + old_attrs_by_id=old_attrs_by_id, + new_attrs_by_id=new_attrs_by_id, + changed_keys=changed_keys, + ) + + def _maintain_views_edge_attrs( + self, + edge_ids: Sequence[int], + attrs: dict[str, Any], + ) -> None: + """ + Bring every registered view up to date after an edge attribute update. + + The edge counterpart of `_maintain_views_node_attrs`. There is no + ``edge_updated`` signal, so this only propagates data — no notification + step and no before/after snapshots. + + `edge_ids` are this graph's own edge ids, which is what views map from. + """ + for view in self._views: + view._apply_root_edge_attrs(edge_ids=edge_ids, attrs=attrs) + + def _maintain_views_attr_key(self, schema: AttrSchema, mode: Literal["node", "edge"]) -> None: + """ + Bring every registered view up to date after a new attribute key is added. + + The schema counterpart of `_maintain_views_node_attrs`, called by concrete + ``add_node_attr_key`` / ``add_edge_attr_key`` implementations once the key + exists on this graph. A view that keeps its own copy of the attributes has + to grow the column too, otherwise it keeps reporting a stale schema and + rejects later writes to the new key. + + Adding a key is a schema operation, so this runs once per key rather than + once per row of a write. + """ + for view in self._views: + view._apply_root_attr_key(schema, mode) + + def _maintain_views_remove_attr_key(self, key: str, mode: Literal["node", "edge"]) -> None: + """ + Bring every registered view up to date after an attribute key is removed. + + The mirror of `_maintain_views_attr_key`, called by concrete + ``remove_node_attr_key`` / ``remove_edge_attr_key`` implementations once + the key is gone from this graph. Without it a view that keeps its own copy + of the schema or of the attributes keeps advertising a column that no + longer exists on the root. + """ + for view in self._views: + view._apply_root_remove_attr_key(key, mode) + @staticmethod def _validate_attributes( attrs: dict[str, Any], diff --git a/src/tracksdata/graph/_graph_view.py b/src/tracksdata/graph/_graph_view.py index 1be0cf55..2a326f1b 100644 --- a/src/tracksdata/graph/_graph_view.py +++ b/src/tracksdata/graph/_graph_view.py @@ -1,5 +1,5 @@ from collections.abc import Callable, Sequence -from typing import Any, Literal, cast, overload +from typing import Any, Literal, overload import bidict import numpy as np @@ -114,6 +114,10 @@ def __init__( self._sync = sync self._out_of_sync = False + # Register with the root so that writes made directly to the root are + # applied to this view. Held weakly, so no explicit teardown is needed. + root._views.add(self) + # Existing for API compatibility for the SQLGraph generating GraphView, # but RXGraph always uses the root graph's attributes and just filtering them self._node_attr_keys = node_attr_keys @@ -305,41 +309,15 @@ def add_node_attr_key( dtype: pl.DataType | None = None, default_value: Any = None, ) -> None: - # Delegate to root with all parameters (root handles overloading) + # Delegate to root with all parameters (root handles overloading). The root + # applies the key back to this view -- and to its sibling views -- through + # `_maintain_views_attr_key`, so there is nothing to do locally here. self._root.add_node_attr_key(key_or_schema, dtype, default_value) - # Extract key for local tracking - if isinstance(key_or_schema, AttrSchema): - key = key_or_schema.key - else: - key = key_or_schema - - if self._node_attr_keys is not None: - self._node_attr_keys.append(key) - - # Sync logic - if not self._is_root_rx_graph: - if self.sync: - # Get the schema from root to get the actual default value used - schema = self._root._node_attr_schemas()[key] - # Apply to local rx_graph - rx_graph = self.rx_graph - for node_id in rx_graph.node_indices(): - rx_graph[node_id][key] = schema.default_value - else: - self._out_of_sync = True - def remove_node_attr_key(self, key: str) -> None: + # See `add_node_attr_key`: the root drops the key from this view -- and from + # its sibling views -- through `_maintain_views_remove_attr_key`. self._root.remove_node_attr_key(key) - if self._node_attr_keys is not None and key in self._node_attr_keys: - self._node_attr_keys.remove(key) - - if not self._is_root_rx_graph: - if self.sync: - for node_id in self.rx_graph.node_indices(): - self.rx_graph[node_id].pop(key, None) - else: - self._out_of_sync = True def add_edge_attr_key( self, @@ -347,40 +325,88 @@ def add_edge_attr_key( dtype: pl.DataType | None = None, default_value: Any = None, ) -> None: - # Delegate to root with all parameters (root handles overloading) + # See `add_node_attr_key`: the root propagates the key back to this view. self._root.add_edge_attr_key(key_or_schema, dtype, default_value) - # Extract key for local tracking - if isinstance(key_or_schema, AttrSchema): - key = key_or_schema.key - else: - key = key_or_schema + def _apply_root_attr_key(self, schema: AttrSchema, mode: Literal["node", "edge"]) -> None: + """ + Absorb a new attribute key registered on the root graph. - if self._edge_attr_keys is not None: - self._edge_attr_keys.append(key) - - # Sync logic - if not self._is_root_rx_graph: - if self.sync: - # Get the schema from root to get the actual default value used - schema = self._root._edge_attr_schemas()[key] - # Apply to local rx_graph - for _, _, edge_attr in self.rx_graph.weighted_edge_list(): - edge_attr[key] = schema.default_value - else: - self._out_of_sync = True + A view that pins an explicit key list has to record the new key there, or + it keeps reporting a stale schema. Beyond that, when the root is a + rustworkx graph the view shares the root's attribute dicts, so the column + already exists on every row; otherwise (e.g. a SQLGraph root) the view + holds its own copy and has to grow the column itself, filling existing + rows with the schema's default value. + + Parameters + ---------- + schema : AttrSchema + The schema of the newly added key, as stored by the root. The default + value is read from here rather than from the caller's arguments, since + the root may have inferred it. + mode : Literal["node", "edge"] + Whether the key was added to the nodes or the edges. + """ + local_keys = self._node_attr_keys if mode == "node" else self._edge_attr_keys + if local_keys is not None and schema.key not in local_keys: + local_keys.append(schema.key) + + if self._is_root_rx_graph: + return + + if not self.sync: + self._out_of_sync = True + return + + rx_graph = self.rx_graph + if mode == "node": + for node_id in rx_graph.node_indices(): + rx_graph[node_id][schema.key] = schema.default_value + else: + for _, _, edge_attr in rx_graph.weighted_edge_list(): + edge_attr[schema.key] = schema.default_value def remove_edge_attr_key(self, key: str) -> None: + # See `remove_node_attr_key`: the root propagates the removal back here. self._root.remove_edge_attr_key(key) - if self._edge_attr_keys is not None and key in self._edge_attr_keys: - self._edge_attr_keys.remove(key) - # because attributes are passed by reference, we need don't need if both are rustworkx graphs - if not self._is_root_rx_graph: - if self.sync: - for edge_attr in self.rx_graph.edges(): - edge_attr.pop(key, None) - else: - self._out_of_sync = True + + def _apply_root_remove_attr_key(self, key: str, mode: Literal["node", "edge"]) -> None: + """ + Absorb an attribute key removed from the root graph. + + The mirror of `_apply_root_attr_key`: a view pinning an explicit key list + has to forget the key there, or it keeps reporting a column that no longer + exists. Beyond that, when the root is a rustworkx graph the view shares the + root's attribute dicts, so the column is already gone from every row; + otherwise (e.g. a SQLGraph root) the view holds its own copy and has to + drop the column itself. + + Parameters + ---------- + key : str + The key that was removed from the root. + mode : Literal["node", "edge"] + Whether the key was removed from the nodes or the edges. + """ + local_keys = self._node_attr_keys if mode == "node" else self._edge_attr_keys + if local_keys is not None and key in local_keys: + local_keys.remove(key) + + if self._is_root_rx_graph: + return + + if not self.sync: + self._out_of_sync = True + return + + rx_graph = self.rx_graph + if mode == "node": + for node_id in rx_graph.node_indices(): + rx_graph[node_id].pop(key, None) + else: + for _, _, edge_attr in rx_graph.weighted_edge_list(): + edge_attr.pop(key, None) def add_node( self, @@ -388,12 +414,11 @@ def add_node( validate_keys: bool = True, index: int | None = None, ) -> int: - with self._root.node_added.blocked(): - parent_node_id = self._root.add_node( - attrs=attrs, - validate_keys=validate_keys, - index=index, - ) + parent_node_id = self._root.add_node( + attrs=attrs, + validate_keys=validate_keys, + index=index, + ) if self.sync: # Local primitive: pure rx_graph + _time_to_nodes, no validation, no signal. @@ -402,14 +427,12 @@ def add_node( else: self._out_of_sync = True - emit_node_added_events(self._root.node_added, [(parent_node_id, attrs)]) emit_node_added_events(self.node_added, [(parent_node_id, attrs)]) return parent_node_id def bulk_add_nodes(self, nodes: list[dict[str, Any]], indices: list[int] | None = None) -> list[int]: - with self._root.node_added.blocked(): - parent_node_ids = self._root.bulk_add_nodes(nodes, indices=indices) + parent_node_ids = self._root.bulk_add_nodes(nodes, indices=indices) if self._is_root_rx_graph: # The rx root stored these exact dict objects by reference (and does not @@ -431,7 +454,6 @@ def bulk_add_nodes(self, nodes: list[dict[str, Any]], indices: list[int] | None else: self._out_of_sync = True - emit_node_added_events(self._root.node_added, zip(parent_node_ids, emitted_nodes, strict=True)) emit_node_added_events(self.node_added, zip(parent_node_ids, emitted_nodes, strict=True)) return parent_node_ids @@ -486,9 +508,9 @@ def bulk_remove_nodes(self, node_ids: Sequence[int]) -> None: raise ValueError(f"Node {missing[0]} does not exist in the graph.") view_signal_on = is_signal_on(self.node_removed) - root_signal_on = is_signal_on(self._root.node_removed) old_attrs_per_node: dict[int, dict[str, Any]] = {} - if view_signal_on or root_signal_on: + if view_signal_on: + # Must be captured before removal, while the attributes still exist. # Single batched query instead of one filter+materialize per node. # include_key defaults to False, so NODE_ID is excluded from each attrs # dict, matching the previous per-node NodeInterface.to_dict() behaviour. @@ -498,8 +520,7 @@ def bulk_remove_nodes(self, node_ids: Sequence[int]) -> None: .rows_by_key(key=DEFAULT_ATTR_KEYS.NODE_ID, named=True, unique=True) ) - with self._root.node_removed.blocked(): - self._root.bulk_remove_nodes(node_ids) + self._root.bulk_remove_nodes(node_ids) if self.sync: local_ids = [self._external_to_local[nid] for nid in node_ids] @@ -514,8 +535,6 @@ def bulk_remove_nodes(self, node_ids: Sequence[int]) -> None: else: self._out_of_sync = True - if root_signal_on: - emit_node_removed_events(self._root.node_removed, ((nid, old_attrs_per_node[nid]) for nid in node_ids)) if view_signal_on: emit_node_removed_events(self.node_removed, ((nid, old_attrs_per_node[nid]) for nid in node_ids)) @@ -967,82 +986,116 @@ def update_node_attrs( attrs: dict[str, Any], node_ids: Sequence[int] | None = None, ) -> None: + """ + Update node attributes through this view. + + Delegates to the root, which applies the change back to this view (and any + sibling views) through the normal maintenance path, so there is a single + implementation of "absorb a node attribute change". + """ if node_ids is None: node_ids = self.node_ids() else: node_ids = list(node_ids) - # Capture signal state once so slots connecting mid-call cannot toggle behavior - # between the old/new attr captures or between the two emit blocks. - view_signal_on = is_signal_on(self.node_updated) - root_signal_on = is_signal_on(self._root.node_updated) - if view_signal_on or root_signal_on: - existing_keys = set(self._root.node_attr_keys(return_ids=True)) - signal_keys = list( - dict.fromkeys( - k - for k in [ - DEFAULT_ATTR_KEYS.NODE_ID, - DEFAULT_ATTR_KEYS.T, - DEFAULT_ATTR_KEYS.Z, - DEFAULT_ATTR_KEYS.Y, - DEFAULT_ATTR_KEYS.X, - DEFAULT_ATTR_KEYS.BBOX, - *attrs.keys(), - ] - if k in existing_keys - ) - ) - old_attrs_by_id = ( - self._root.filter(node_ids=node_ids) - .node_attrs(attr_keys=signal_keys) - .rows_by_key(key=DEFAULT_ATTR_KEYS.NODE_ID, named=True, unique=True, include_key=True) - ) + self._root.update_node_attrs(node_ids=node_ids, attrs=attrs) - # Block root signal so it doesn't fire while the view is still in old state; - # re-emit at the end after both root and view are consistent. - with self._root.node_updated.blocked(): - self._root.update_node_attrs( - node_ids=node_ids, - attrs=attrs, - ) - # because attributes are passed by reference, we need don't need if both are rustworkx graphs - if not self._is_root_rx_graph: - if self.sync: + def _needs_root_node_attrs(self) -> tuple[bool, bool]: + """ + Whether this view needs before/after snapshots of a root node update. + + Queried by the root (see `BaseGraph._views_need_node_attrs`) so it can skip + building snapshots no view will read. + + Returns + ------- + tuple[bool, bool] + ``(needs_old, needs_new)``. Old values are only used to emit + ``node_updated``. New values are needed on top of that when this view + keeps its own copy of the attributes and has to be written through -- + which an out-of-sync view does not do either, it only marks itself + stale. + """ + listening = is_signal_on(self.node_updated) + writes_through = self.sync and not self._is_root_rx_graph + return listening, listening or writes_through + + def _apply_root_node_attrs( + self, + *, + node_ids: Sequence[int], + old_attrs_by_id: dict[int, dict[str, Any]] | None, + new_attrs_by_id: dict[int, dict[str, Any]] | None, + changed_keys: set[str], + ) -> None: + """ + Absorb a node attribute update made directly on the root graph. + + Brings this view up to date and *then* emits its own ``node_updated`` + signal for the nodes it contains. Keeping the view consistent with its + root is an invariant, so the local update happens whether or not + anything is listening — only the emission is conditional. + + Parameters + ---------- + node_ids : Sequence[int] + Nodes updated on the root, in root ids. + old_attrs_by_id : dict[int, dict[str, Any]] | None + Attributes before the update, keyed by root id. ``None`` when no view + asked for them, in which case nothing here emits a signal. + new_attrs_by_id : dict[int, dict[str, Any]] | None + Attributes after the update, keyed by root id. ``None`` when no view + asked for them. + changed_keys : set[str] + Attribute keys written by this update. + """ + listening = is_signal_on(self.node_updated) + + # Bail out before scanning `node_ids`, which is one membership test per + # updated node. A view sharing the root's attribute dicts is already + # current, so with nothing listening there is no work at all. + if self._is_root_rx_graph and not listening: + return + + # An out-of-sync view is only marked stale, so it needs no scan either. + if not self._is_root_rx_graph and not self.sync: + self._out_of_sync = True + if not listening: + return + + in_view = [node_id for node_id in node_ids if self.has_node(node_id)] + if not in_view: + return + + # Maintain first. When the root is a rustworkx graph the view shares the + # root's attribute dicts, so the values are already current and writing + # again would be redundant. Otherwise (e.g. a SQLGraph root) the view + # holds its own copy and has to be written through. + if not self._is_root_rx_graph and self.sync: + # A view may track only a subset of the root's keys; the ones it + # left out have no local column to write to, so they are skipped + # rather than forwarded (the local store would reject them). + local_keys = set(self.node_attr_keys(return_ids=True)) + local_attrs = { + key: [new_attrs_by_id[node_id][key] for node_id in in_view] + for key in changed_keys + if key in local_keys and all(key in new_attrs_by_id[node_id] for node_id in in_view) + } + if local_attrs: with self.node_updated.blocked(): - super().update_node_attrs( - node_ids=self._map_to_local(node_ids), - attrs=attrs, + RustWorkXGraph.update_node_attrs( + self, + node_ids=self._map_to_local(in_view), + attrs=local_attrs, ) - else: - self._out_of_sync = True - - if view_signal_on or root_signal_on: - old_attrs_by_id = cast(dict[int, dict[str, Any]], old_attrs_by_id) # for mypy - # Derive new_attrs by overlaying applied `attrs` onto old_attrs, instead of - # re-querying root. Mirrors the broadcasting semantics of - # `_root.update_node_attrs`: scalars apply to all nodes, sequences index by - # position in `node_ids`. - new_attrs_by_id: dict[int, dict[str, Any]] = {} - for i, node_id in enumerate(node_ids): - new_attrs = dict(old_attrs_by_id[node_id]) - for k, v in attrs.items(): - if k in new_attrs: - new_attrs[k] = v if np.isscalar(v) else v[i] - new_attrs_by_id[node_id] = new_attrs - changed_keys = set(attrs.keys()) - if root_signal_on: - emit_node_updated_events( - self._root.node_updated, - ((node_id, old_attrs_by_id[node_id], new_attrs_by_id[node_id]) for node_id in node_ids), - changed_keys, - ) - if view_signal_on: - emit_node_updated_events( - self.node_updated, - ((node_id, old_attrs_by_id[node_id], new_attrs_by_id[node_id]) for node_id in node_ids), - changed_keys, - ) + + # Notify second, now that root and view agree. + if listening: + emit_node_updated_events( + self.node_updated, + ((node_id, old_attrs_by_id[node_id], new_attrs_by_id[node_id]) for node_id in in_view), + changed_keys, + ) def update_edge_attrs( self, @@ -1050,6 +1103,12 @@ def update_edge_attrs( attrs: dict[str, Any], edge_ids: Sequence[int] | None = None, ) -> None: + """ + Update edge attributes through this view. + + Delegates to the root, which applies the change back to this view (and any + sibling views) through the normal maintenance path. + """ if edge_ids is None: edge_ids = self.edge_ids() @@ -1057,15 +1116,55 @@ def update_edge_attrs( edge_ids=edge_ids, attrs=attrs, ) - # because attributes are passed by reference, we need don't need if both are rustworkx graphs - if not self._is_root_rx_graph: - if self.sync: - super().update_edge_attrs( - edge_ids=[self._edge_map_from_root[eid] for eid in edge_ids], - attrs=attrs, - ) - else: - self._out_of_sync = True + + def _apply_root_edge_attrs( + self, + *, + edge_ids: Sequence[int], + attrs: dict[str, Any], + ) -> None: + """ + Absorb an edge attribute update made directly on the root graph. + + Parameters + ---------- + edge_ids : Sequence[int] + Edges updated on the root, in root edge ids. + attrs : dict[str, Any] + The attributes written, in the same form passed to + ``update_edge_attrs``. + """ + # When the root is a rustworkx graph the view shares the root's attribute + # dicts, so the values are already current. + if self._is_root_rx_graph: + return + + # Keep positions, not just ids: per-edge values are positional, so the + # in-view subset has to be selected by the same indices. + in_view = [(i, edge_id) for i, edge_id in enumerate(edge_ids) if edge_id in self._edge_map_from_root] + if not in_view: + return + + if not self.sync: + self._out_of_sync = True + return + + # See `_apply_root_node_attrs`: keys this view does not track have no + # local column and are skipped. + local_keys = set(self.edge_attr_keys(return_ids=True)) + positions = [i for i, _ in in_view] + local_attrs = { + key: value if np.isscalar(value) else [value[i] for i in positions] + for key, value in attrs.items() + if key in local_keys + } + if not local_attrs: + return + + super().update_edge_attrs( + edge_ids=[self._edge_map_from_root[edge_id] for _, edge_id in in_view], + attrs=local_attrs, + ) def in_degree(self, node_ids: list[int] | int | None = None) -> list[int] | int: """ diff --git a/src/tracksdata/graph/_mapped_graph_mixin.py b/src/tracksdata/graph/_mapped_graph_mixin.py index 19d4f1fa..a7d1cb73 100644 --- a/src/tracksdata/graph/_mapped_graph_mixin.py +++ b/src/tracksdata/graph/_mapped_graph_mixin.py @@ -52,12 +52,20 @@ def __init__(self, id_map: dict[int, int] | bidict.bidict[int, int] | None = Non self._external_to_local = self._local_to_external.inverse def __getstate__(self) -> dict[str, Any]: - data = self.__dict__.copy() + # Defer to the next class in the MRO (BaseGraph, in every concrete use) + # so its own exclusions are honoured; fall back to __dict__ if this + # mixin is ever used without such a base. + parent = getattr(super(), "__getstate__", None) + data = parent() if parent is not None else self.__dict__.copy() del data["_external_to_local"] return data def __setstate__(self, state: dict[str, Any]) -> None: - self.__dict__.update(state) + parent = getattr(super(), "__setstate__", None) + if parent is not None: + parent(state) + else: + self.__dict__.update(state) self._external_to_local = self._local_to_external.inverse @overload diff --git a/src/tracksdata/graph/_rustworkx_graph.py b/src/tracksdata/graph/_rustworkx_graph.py index d7466f08..7c637458 100644 --- a/src/tracksdata/graph/_rustworkx_graph.py +++ b/src/tracksdata/graph/_rustworkx_graph.py @@ -1,6 +1,7 @@ import operator from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any +from weakref import WeakSet import bidict import numpy as np @@ -1089,6 +1090,8 @@ def add_node_attr_key( # Store schema self.__node_attr_schemas[schema.key] = schema + self._maintain_views_attr_key(schema, "node") + def remove_node_attr_key(self, key: str) -> None: """ Remove an existing node attribute key from the graph. @@ -1103,6 +1106,8 @@ def remove_node_attr_key(self, key: str) -> None: for node_attr in self.rx_graph.nodes(): node_attr.pop(key, None) + self._maintain_views_remove_attr_key(key, "node") + def add_edge_attr_key( self, key_or_schema: str | AttrSchema, @@ -1129,6 +1134,8 @@ def add_edge_attr_key( # Store schema self.__edge_attr_schemas[schema.key] = schema + self._maintain_views_attr_key(schema, "edge") + def remove_edge_attr_key(self, key: str) -> None: """ Remove an existing edge attribute key from the graph. @@ -1140,6 +1147,8 @@ def remove_edge_attr_key(self, key: str) -> None: for edge_attr in self.rx_graph.edges(): edge_attr.pop(key, None) + self._maintain_views_remove_attr_key(key, "edge") + def _node_attrs_from_node_ids( self, *, @@ -1313,7 +1322,15 @@ def update_node_attrs( else: node_ids = list(node_ids) - if is_signal_on(self.node_updated): + # Views must be maintained even with no listeners, but only some of them + # read the before/after snapshots -- see `_views_need_node_attrs`. + signal_on = is_signal_on(self.node_updated) + views_need_old, views_need_new = self._views_need_node_attrs() + needs_old = signal_on or views_need_old + needs_new = signal_on or views_need_new + old_attrs_by_id = None + new_attrs_by_id = None + if needs_old: old_attrs_by_id = {node_id: dict(self._graph[node_id]) for node_id in node_ids} for key, value in attrs.items(): @@ -1330,11 +1347,21 @@ def update_node_attrs( for node_id, v in zip(node_ids, value, strict=False): self._graph[node_id][key] = v - if is_signal_on(self.node_updated): + changed_keys = set(attrs.keys()) + if needs_new: + new_attrs_by_id = {node_id: dict(self._graph[node_id]) for node_id in node_ids} + if signal_on: emit_node_updated_events( self.node_updated, - ((node_id, old_attrs_by_id[node_id], dict(self._graph[node_id])) for node_id in node_ids), - set(attrs.keys()), + ((node_id, old_attrs_by_id[node_id], new_attrs_by_id[node_id]) for node_id in node_ids), + changed_keys, + ) + if self._views: + self._maintain_views_node_attrs( + node_ids=node_ids, + old_attrs_by_id=old_attrs_by_id, + new_attrs_by_id=new_attrs_by_id, + changed_keys=changed_keys, ) def update_edge_attrs( @@ -1379,6 +1406,9 @@ def update_edge_attrs( for key, value in broadcast_attrs.items(): edge_attr[key] = value[i] + if self._views: + self._maintain_views_edge_attrs(edge_ids=edge_ids, attrs=broadcast_attrs) + def assign_tracklet_ids( self, output_key: str = DEFAULT_ATTR_KEYS.TRACKLET_ID, @@ -1999,23 +2029,52 @@ def update_node_attrs( external_node_ids = self.node_ids() if node_ids is None else node_ids local_node_ids = self._map_to_local(external_node_ids) - if is_signal_on(self.node_updated): + # Views must be maintained even with no listeners, but only some of them + # read the before/after snapshots -- see `_views_need_node_attrs`. + signal_on = is_signal_on(self.node_updated) + views_need_old, views_need_new = self._views_need_node_attrs() + needs_old = signal_on or views_need_old + needs_new = signal_on or views_need_new + old_attrs_by_id = None + new_attrs_by_id = None + if needs_old: old_attrs_by_id = { external_node_id: dict(self._graph[local_node_id]) for external_node_id, local_node_id in zip(external_node_ids, local_node_ids, strict=True) } - with self.node_updated.blocked(): - super().update_node_attrs(attrs=attrs, node_ids=local_node_ids) - - if is_signal_on(self.node_updated) and old_attrs_by_id is not None: + # Suppress both the signal and view maintenance during the super() call: + # it works in local ids, whereas views map from this graph's external + # ids. Both are re-done below with external ids. + saved_views = self._views + self._views = WeakSet() + try: + with self.node_updated.blocked(): + super().update_node_attrs(attrs=attrs, node_ids=local_node_ids) + finally: + self._views = saved_views + + changed_keys = set(attrs.keys()) + if needs_new: + new_attrs_by_id = { + external_node_id: dict(self._graph[local_node_id]) + for external_node_id, local_node_id in zip(external_node_ids, local_node_ids, strict=True) + } + if signal_on: emit_node_updated_events( self.node_updated, ( - (external_node_id, old_attrs_by_id[external_node_id], dict(self._graph[local_node_id])) - for external_node_id, local_node_id in zip(external_node_ids, local_node_ids, strict=True) + (external_node_id, old_attrs_by_id[external_node_id], new_attrs_by_id[external_node_id]) + for external_node_id in external_node_ids ), - set(attrs.keys()), + changed_keys, + ) + if self._views: + self._maintain_views_node_attrs( + node_ids=external_node_ids, + old_attrs_by_id=old_attrs_by_id, + new_attrs_by_id=new_attrs_by_id, + changed_keys=changed_keys, ) def bulk_remove_nodes(self, node_ids: Sequence[int]) -> None: diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index f3602e86..8ee3d817 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/src/tracksdata/graph/_sql_graph.py @@ -653,6 +653,7 @@ def __init__( engine_kwargs: dict[str, Any] | None = None, overwrite: bool = False, ): + super().__init__() self._url = sa.engine.URL.create( drivername, username=username, @@ -1971,6 +1972,8 @@ def add_node_attr_key( node_schemas[schema.key] = schema self.__node_attr_schemas = node_schemas + self._maintain_views_attr_key(schema, "node") + def remove_node_attr_key(self, key: str) -> None: if key not in self.node_attr_keys(): raise ValueError(f"Node attribute key {key} does not exist") @@ -1983,6 +1986,8 @@ def remove_node_attr_key(self, key: str) -> None: node_schemas.pop(key, None) self.__node_attr_schemas = node_schemas + self._maintain_views_remove_attr_key(key, "node") + def add_edge_attr_key( self, key_or_schema: str | AttrSchema, @@ -1998,6 +2003,8 @@ def add_edge_attr_key( edge_schemas[schema.key] = schema self.__edge_attr_schemas = edge_schemas + self._maintain_views_attr_key(schema, "edge") + def remove_edge_attr_key(self, key: str) -> None: if key not in self.edge_attr_keys(): raise ValueError(f"Edge attribute key {key} does not exist") @@ -2007,6 +2014,8 @@ def remove_edge_attr_key(self, key: str) -> None: edge_schemas.pop(key, None) self.__edge_attr_schemas = edge_schemas + self._maintain_views_remove_attr_key(key, "edge") + def num_edges(self) -> int: with Session(self._engine) as session: return int(session.query(self.Edge).count()) @@ -2187,7 +2196,16 @@ def update_node_attrs( return attr_keys = self.node_attr_keys() - if is_signal_on(self.node_updated): + # Views must be maintained even with no listeners, but only some of them + # read the before/after snapshots -- see `_views_need_node_attrs`. Each + # snapshot is a full query over the updated rows, so skipping one matters. + signal_on = is_signal_on(self.node_updated) + views_need_old, views_need_new = self._views_need_node_attrs() + needs_old = signal_on or views_need_old + needs_new = signal_on or views_need_new + old_attrs_by_id = None + new_attrs_by_id = None + if needs_old: old_df = self.filter(node_ids=updated_node_ids).node_attrs( attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, *attr_keys] ) @@ -2197,17 +2215,31 @@ def update_node_attrs( self._update_table(self.Node, node_ids, DEFAULT_ATTR_KEYS.NODE_ID, attrs) - if is_signal_on(self.node_updated): + changed_keys = set(attrs.keys()) + if needs_new: + # `needs_old` is exactly "somebody will emit", and an emitted payload has + # to carry every attribute. Nobody emitting means the snapshot is only + # feeding write-through into views, which reads the changed keys alone, + # so the query can skip the rest of the columns. + new_attr_keys = attr_keys if needs_old else [k for k in attr_keys if k in changed_keys] new_df = self.filter(node_ids=updated_node_ids).node_attrs( - attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, *attr_keys] + attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, *new_attr_keys] ) new_attrs_by_id = new_df.rows_by_key( key=DEFAULT_ATTR_KEYS.NODE_ID, named=True, unique=True, include_key=True ) + if signal_on: emit_node_updated_events( self.node_updated, ((node_id, old_attrs_by_id[node_id], new_attrs_by_id[node_id]) for node_id in updated_node_ids), - set(attrs.keys()), + changed_keys, + ) + if self._views: + self._maintain_views_node_attrs( + node_ids=updated_node_ids, + old_attrs_by_id=old_attrs_by_id, + new_attrs_by_id=new_attrs_by_id, + changed_keys=changed_keys, ) def update_edge_attrs( @@ -2218,6 +2250,10 @@ def update_edge_attrs( ) -> None: self._update_table(self.Edge, edge_ids, DEFAULT_ATTR_KEYS.EDGE_ID, attrs) + if self._views: + updated_edge_ids = self.edge_ids() if edge_ids is None else list(edge_ids) + self._maintain_views_edge_attrs(edge_ids=updated_edge_ids, attrs=attrs) + def assign_tracklet_ids( self, output_key: str = DEFAULT_ATTR_KEYS.TRACKLET_ID, @@ -2485,13 +2521,13 @@ def _sqlite_table_dump( _drop_scratch_table(source_root._engine, selected) def __getstate__(self) -> dict: - data_dict = self.__dict__.copy() + data_dict = super().__getstate__() for k in ["Base", "Node", "Edge", "Overlap", "Metadata", "_engine"]: del data_dict[k] return data_dict def __setstate__(self, state: dict) -> None: - self.__dict__.update(state) + super().__setstate__(state) # recreate deleted objects self._engine = sa.create_engine(self._url, **self._engine_kwargs) self._define_schema(overwrite=False) diff --git a/src/tracksdata/graph/_test/test_graph_view_signals.py b/src/tracksdata/graph/_test/test_graph_view_signals.py index 311288fc..bec5cc94 100644 --- a/src/tracksdata/graph/_test/test_graph_view_signals.py +++ b/src/tracksdata/graph/_test/test_graph_view_signals.py @@ -4,21 +4,33 @@ either the root or the view must see the two graphs in a consistent state. A listener attached to root that queries the view (or vice versa) must not observe ghost or stale nodes. + +The second half of this module covers the *reverse* direction: a GraphView is +meant to be a live view, so writes made directly to the **root** must be +visible in the view and must re-emit ``view.node_updated`` for the nodes the +view contains. Edge attributes propagate the same way, minus the notification — +there is no ``edge_updated`` signal. """ +import gc +import pickle + import polars as pl +import pytest +from tracksdata.attrs import NodeAttr from tracksdata.constants import DEFAULT_ATTR_KEYS -from tracksdata.graph import BaseGraph +from tracksdata.graph import BaseGraph, RustWorkXGraph +from tracksdata.graph._rustworkx_graph import IndexedRXGraph -def test_view_node_signals_fire_with_consistent_state(graph_backend: BaseGraph) -> None: - """add_node / remove_node: when either signal fires (on root or view), the - two graphs must agree on `has_node`. +def test_node_signals_fire_after_the_emitting_graph_is_updated(graph_backend: BaseGraph) -> None: + """add_node / remove_node: a signal must reflect the graph that emitted it. - Today used to fail on `remove_node` because `GraphView.remove_node` does not - block the root signal — root emits while the view's local rx_graph still - holds the node. + Each graph is responsible for its own signal only. A listener on the root + sees the root updated; a listener on the view sees the view updated. Neither + is required to observe the *other* graph in any particular state — see + royerlab/tracksdata#324. """ graph_backend.add_node_attr_key("x", pl.Float64) graph_backend.add_node({"t": 0, "x": 0.0}) @@ -26,36 +38,42 @@ def test_view_node_signals_fire_with_consistent_state(graph_backend: BaseGraph) view = graph_backend.filter().subgraph() observations: list = [] - def make_slot(source: str, signal: str): + def make_slot(graph: BaseGraph, source: str, signal: str): def slot(node_ids: list[int], *_args) -> None: for node_id in node_ids: - observations.append((source, signal, node_id, graph_backend.has_node(node_id), view.has_node(node_id))) + observations.append((source, signal, node_id, graph.has_node(node_id))) return slot - graph_backend.node_added.connect(make_slot("root", "added")) - graph_backend.node_removed.connect(make_slot("root", "removed")) - view.node_added.connect(make_slot("view", "added")) - view.node_removed.connect(make_slot("view", "removed")) + graph_backend.node_added.connect(make_slot(graph_backend, "root", "added")) + graph_backend.node_removed.connect(make_slot(graph_backend, "root", "removed")) + view.node_added.connect(make_slot(view, "view", "added")) + view.node_removed.connect(make_slot(view, "view", "removed")) new_id = view.add_node({"t": 1, "x": 1.0}) view.remove_node(new_id) - inconsistent = [obs for obs in observations if obs[3] != obs[4]] + # every "added" must see the node present, every "removed" must see it gone, + # each in the graph that emitted the signal + wrong = [obs for obs in observations if obs[3] != (obs[1] == "added")] detail = "\n".join( - f" {source}.{signal}(node={nid}): root.has_node={rh}, view.has_node={vh}" - for source, signal, nid, rh, vh in inconsistent + f" {source}.{signal}(node={nid}): {source}.has_node={present}" for source, signal, nid, present in wrong ) - assert not inconsistent, f"Listener saw root and view in inconsistent state at signal time:\n{detail}" + assert not wrong, f"Signal did not reflect the state of the graph that emitted it:\n{detail}" + # both graphs emitted both events + assert {(source, signal) for source, signal, _, _ in observations} == { + ("root", "added"), + ("root", "removed"), + ("view", "added"), + ("view", "removed"), + } -def test_view_update_node_attrs_signal_fires_with_consistent_value(graph_backend: BaseGraph) -> None: - """update_node_attrs: when either signal fires, root and view must hold - the same value for the updated attribute. +def test_update_node_attrs_signal_reflects_the_emitting_graph(graph_backend: BaseGraph) -> None: + """update_node_attrs: each graph's signal must carry that graph's new value. - This used to fail on backends where root and view do not share an attribute - storage (SQLGraph): root emits with the new value while the view's local - rx_graph still holds the old one. + As with add/remove, a listener is only promised that the graph it subscribed + to is current — not that root and view agree at that instant. """ graph_backend.add_node_attr_key("x", pl.Float64) node_id = graph_backend.add_node({"t": 0, "x": 0.0}) @@ -67,20 +85,407 @@ def attr_value(graph: BaseGraph, nid: int) -> float: df = graph.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, "x"]) return df.filter(pl.col(DEFAULT_ATTR_KEYS.NODE_ID) == nid)["x"].item() - def make_slot(source: str): - def slot(node_ids: list[int], _old: list[dict], _new: list[dict]) -> None: + def make_slot(graph: BaseGraph, source: str): + def slot(node_ids: list[int], _old: list[dict], _new: list[dict], *_rest) -> None: for nid in node_ids: - observations.append((source, nid, attr_value(graph_backend, nid), attr_value(view, nid))) + observations.append((source, nid, attr_value(graph, nid))) return slot - graph_backend.node_updated.connect(make_slot("root")) - view.node_updated.connect(make_slot("view")) + graph_backend.node_updated.connect(make_slot(graph_backend, "root")) + view.node_updated.connect(make_slot(view, "view")) view.update_node_attrs(attrs={"x": 5.0}, node_ids=[node_id]) - inconsistent = [obs for obs in observations if obs[2] != obs[3]] - detail = "\n".join( - f" {source}.node_updated(node={nid}): root.x={rx}, view.x={vx}" for source, nid, rx, vx in inconsistent - ) - assert not inconsistent, f"Listener saw root and view holding different attribute values at signal time:\n{detail}" + stale = [obs for obs in observations if obs[2] != 5.0] + detail = "\n".join(f" {source}.node_updated(node={nid}): {source}.x={value}" for source, nid, value in stale) + assert not stale, f"Signal fired before the emitting graph held the new value:\n{detail}" + assert {source for source, _, _ in observations} == {"root", "view"} + + +# -------------------------------------------------------------------------- +# Root -> view propagation ("the view is a live view") +# -------------------------------------------------------------------------- + + +class _UpdateRecorder: + """Collects every ``node_updated`` emission from a graph. + + Holds a reference to its own bound slot so the connection is not dropped + by signal implementations that keep only weak references to listeners. + """ + + def __init__(self, graph: BaseGraph) -> None: + self.calls: list[tuple[list[int], list[dict], list[dict], set[str]]] = [] + graph.node_updated.connect(self._slot) + + def _slot(self, node_ids: list[int], old: list[dict], new: list[dict], changed_keys: set[str]) -> None: + self.calls.append((list(node_ids), old, new, changed_keys)) + + def __len__(self) -> int: + return len(self.calls) + + +def _record_updates(graph: BaseGraph) -> _UpdateRecorder: + """Connect to ``graph.node_updated`` and collect every emission.""" + return _UpdateRecorder(graph) + + +def _value_of(graph: BaseGraph, node_id: int, key: str = "x") -> float: + df = graph.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, key]) + return df.filter(pl.col(DEFAULT_ATTR_KEYS.NODE_ID) == node_id)[key].item() + + +def test_root_update_is_visible_in_view(graph_backend: BaseGraph) -> None: + """A write directly to the root must be readable through the view.""" + graph_backend.add_node_attr_key("x", pl.Float64) + node_id = graph_backend.add_node({"t": 0, "x": 0.0}) + + view = graph_backend.filter().subgraph() + + graph_backend.update_node_attrs(attrs={"x": 7.0}, node_ids=[node_id]) + + assert _value_of(graph_backend, node_id) == 7.0 + assert _value_of(view, node_id) == 7.0, "view returned a stale value after a root write" + + +def test_root_update_emits_view_signal(graph_backend: BaseGraph) -> None: + """A write to the root must re-emit ``node_updated`` on the view.""" + graph_backend.add_node_attr_key("x", pl.Float64) + node_id = graph_backend.add_node({"t": 0, "x": 0.0}) + + view = graph_backend.filter().subgraph() + view_calls = _record_updates(view) + + graph_backend.update_node_attrs(attrs={"x": 3.5}, node_ids=[node_id]) + + assert len(view_calls) == 1, f"expected exactly one view emission, got {len(view_calls)}" + node_ids, old, new, changed_keys = view_calls.calls[0] + assert node_ids == [node_id] + assert old[0]["x"] == 0.0 + assert new[0]["x"] == 3.5 + assert changed_keys == {"x"} + + +def test_root_update_reports_view_node_ids(graph_backend: BaseGraph) -> None: + """Forwarded events must use the IDs the view exposes, not internal ones. + + Relevant for ``IndexedRXGraph`` roots, where the graph's external node IDs + differ from the underlying rustworkx indices. + """ + graph_backend.add_node_attr_key("x", pl.Float64) + node_ids = [graph_backend.add_node({"t": t, "x": 0.0}) for t in range(3)] + + view = graph_backend.filter().subgraph() + view_calls = _record_updates(view) + + graph_backend.update_node_attrs(attrs={"x": 1.0}, node_ids=[node_ids[1]]) + + assert len(view_calls) == 1 + reported = view_calls.calls[0][0] + assert reported == [node_ids[1]] + # the reported id must be addressable on the view + assert view.has_node(reported[0]) + assert _value_of(view, reported[0]) == 1.0 + + +def test_root_update_outside_view_does_not_emit(graph_backend: BaseGraph) -> None: + """Updating a node absent from the view must not emit on the view.""" + graph_backend.add_node_attr_key("x", pl.Float64) + inside = graph_backend.add_node({"t": 0, "x": 0.0}) + outside = graph_backend.add_node({"t": 5, "x": 0.0}) + + view = graph_backend.filter(node_ids=[inside]).subgraph() + assert not view.has_node(outside) + + view_calls = _record_updates(view) + + graph_backend.update_node_attrs(attrs={"x": 9.0}, node_ids=[outside]) + + assert view_calls.calls == [], "view emitted for a node it does not contain" + + +def test_root_update_mixed_batch_filters_to_view_nodes(graph_backend: BaseGraph) -> None: + """A batch spanning in- and out-of-view nodes emits only the in-view subset.""" + graph_backend.add_node_attr_key("x", pl.Float64) + inside = graph_backend.add_node({"t": 0, "x": 0.0}) + outside = graph_backend.add_node({"t": 5, "x": 0.0}) + + view = graph_backend.filter(node_ids=[inside]).subgraph() + view_calls = _record_updates(view) + + graph_backend.update_node_attrs(attrs={"x": 4.0}, node_ids=[inside, outside]) + + assert len(view_calls) == 1 + node_ids, _old, new, _keys = view_calls.calls[0] + assert node_ids == [inside] + assert new[0]["x"] == 4.0 + # both nodes were still written on the root + assert _value_of(graph_backend, outside) == 4.0 + + +def test_root_update_sequence_values_map_per_node(graph_backend: BaseGraph) -> None: + """Per-node sequence values must be reported against the matching node.""" + graph_backend.add_node_attr_key("x", pl.Float64) + first = graph_backend.add_node({"t": 0, "x": 0.0}) + second = graph_backend.add_node({"t": 1, "x": 0.0}) + + view = graph_backend.filter().subgraph() + view_calls = _record_updates(view) + + graph_backend.update_node_attrs(attrs={"x": [11.0, 22.0]}, node_ids=[first, second]) + + assert len(view_calls) == 1 + node_ids, _old, new, _keys = view_calls.calls[0] + by_id = dict(zip(node_ids, [n["x"] for n in new], strict=True)) + assert by_id == {first: 11.0, second: 22.0} + assert _value_of(view, first) == 11.0 + assert _value_of(view, second) == 22.0 + + +def test_root_update_notifies_multiple_sibling_views(graph_backend: BaseGraph) -> None: + """Every registered view containing the node must be notified.""" + graph_backend.add_node_attr_key("x", pl.Float64) + shared = graph_backend.add_node({"t": 0, "x": 0.0}) + other = graph_backend.add_node({"t": 5, "x": 0.0}) + + view_a = graph_backend.filter().subgraph() + view_b = graph_backend.filter(node_ids=[shared]).subgraph() + view_c = graph_backend.filter(node_ids=[other]).subgraph() + + calls_a = _record_updates(view_a) + calls_b = _record_updates(view_b) + calls_c = _record_updates(view_c) + + graph_backend.update_node_attrs(attrs={"x": 6.0}, node_ids=[shared]) + + assert len(calls_a) == 1 + assert len(calls_b) == 1 + assert calls_c.calls == [], "view without the node should not be notified" + + +def test_view_write_emits_exactly_once_on_each_graph(graph_backend: BaseGraph) -> None: + """Writing through the view must not double-emit via the root registry. + + The view unregisters itself from the root while delegating the write, so + the event is emitted once by the view itself rather than twice. + """ + graph_backend.add_node_attr_key("x", pl.Float64) + node_id = graph_backend.add_node({"t": 0, "x": 0.0}) + + view = graph_backend.filter().subgraph() + root_calls = _record_updates(graph_backend) + view_calls = _record_updates(view) + + view.update_node_attrs(attrs={"x": 2.0}, node_ids=[node_id]) + + assert len(root_calls) == 1, f"root emitted {len(root_calls)} times, expected 1" + assert len(view_calls) == 1, f"view emitted {len(view_calls)} times, expected 1" + assert view_calls.calls[0][2][0]["x"] == 2.0 + + +def test_view_is_reregistered_after_writing_through_it(graph_backend: BaseGraph) -> None: + """A write through the view must not permanently unregister it.""" + graph_backend.add_node_attr_key("x", pl.Float64) + node_id = graph_backend.add_node({"t": 0, "x": 0.0}) + + view = graph_backend.filter().subgraph() + + # write through the view first, which temporarily unregisters it + view.update_node_attrs(attrs={"x": 1.0}, node_ids=[node_id]) + + view_calls = _record_updates(view) + # a subsequent *root* write must still reach the view + graph_backend.update_node_attrs(attrs={"x": 8.0}, node_ids=[node_id]) + + assert len(view_calls) == 1, "view stopped receiving root updates after writing through it" + assert view_calls.calls[0][2][0]["x"] == 8.0 + + +def test_root_update_with_no_view_listener_is_harmless(graph_backend: BaseGraph) -> None: + """Forwarding must be a no-op when nothing is connected to the view.""" + graph_backend.add_node_attr_key("x", pl.Float64) + node_id = graph_backend.add_node({"t": 0, "x": 0.0}) + + view = graph_backend.filter().subgraph() + + graph_backend.update_node_attrs(attrs={"x": 5.0}, node_ids=[node_id]) + + assert _value_of(view, node_id) == 5.0 + + +def test_nested_view_receives_root_updates(graph_backend: BaseGraph) -> None: + """A subgraph taken from a view must also observe root writes.""" + graph_backend.add_node_attr_key("x", pl.Float64) + node_id = graph_backend.add_node({"t": 0, "x": 0.0}) + + outer = graph_backend.filter().subgraph() + inner = outer.filter().subgraph() + + outer_calls = _record_updates(outer) + inner_calls = _record_updates(inner) + + graph_backend.update_node_attrs(attrs={"x": 4.5}, node_ids=[node_id]) + + assert len(outer_calls) == 1 + assert len(inner_calls) == 1, "nested view did not receive the root update" + assert _value_of(inner, node_id) == 4.5 + + +def test_unreferenced_view_is_released(graph_backend: BaseGraph) -> None: + """Dropping the last reference to a view must unregister it. + + The registry holds views weakly, so an unreachable view stops costing + anything on subsequent root updates. + """ + graph_backend.add_node_attr_key("x", pl.Float64) + node_id = graph_backend.add_node({"t": 0, "x": 0.0}) + + view = graph_backend.filter().subgraph() + assert len(graph_backend._views) == 1 + + del view + gc.collect() + + assert len(graph_backend._views) == 0, "view was not released from the registry" + + # updating with a released view must not raise + graph_backend.update_node_attrs(attrs={"x": 1.0}, node_ids=[node_id]) + assert _value_of(graph_backend, node_id) == 1.0 + + +def test_referenced_view_is_kept_registered(graph_backend: BaseGraph) -> None: + """A view that is still referenced must survive garbage collection.""" + graph_backend.add_node_attr_key("x", pl.Float64) + node_id = graph_backend.add_node({"t": 0, "x": 0.0}) + + view = graph_backend.filter().subgraph() + gc.collect() + + assert len(graph_backend._views) == 1, "referenced view was dropped" + + graph_backend.update_node_attrs(attrs={"x": 2.0}, node_ids=[node_id]) + assert _value_of(view, node_id) == 2.0 + + +@pytest.mark.parametrize("graph_class", [RustWorkXGraph, IndexedRXGraph]) +def test_view_registry_survives_pickling(graph_class: type[BaseGraph]) -> None: + """Pickling a root or a view must not fail on the weak view registry.""" + graph = graph_class() + graph.add_node_attr_key("x", pl.Float64) + node_id = graph.add_node({"t": 0, "x": 0.0}) + view = graph.filter().subgraph() + + restored_root = pickle.loads(pickle.dumps(graph)) + assert restored_root.has_node(node_id) + # a restored root must still be usable for updates + restored_root.update_node_attrs(attrs={"x": 2.0}, node_ids=[node_id]) + assert _value_of(restored_root, node_id) == 2.0 + + restored_view = pickle.loads(pickle.dumps(view)) + assert restored_view.has_node(node_id) + + +# -------------------------------------------------------------------------- +# Root -> view propagation for edge attributes +# -------------------------------------------------------------------------- + + +def _edge_value_of(graph: BaseGraph, edge_id: int, key: str = "w") -> float: + df = graph.edge_attrs(attr_keys=[DEFAULT_ATTR_KEYS.EDGE_ID, key]) + return df.filter(pl.col(DEFAULT_ATTR_KEYS.EDGE_ID) == edge_id)[key].item() + + +def _graph_with_edge(graph: BaseGraph) -> tuple[int, int]: + """Add two nodes and an edge between them; return (edge_id, other_edge_id).""" + graph.add_node_attr_key("x", pl.Float64) + graph.add_edge_attr_key("w", pl.Float64) + first = graph.add_node({"t": 0, "x": 0.0}) + second = graph.add_node({"t": 1, "x": 1.0}) + third = graph.add_node({"t": 2, "x": 2.0}) + edge_id = graph.add_edge(first, second, {"w": 0.0}) + other_id = graph.add_edge(second, third, {"w": 0.0}) + return edge_id, other_id + + +def test_root_edge_update_is_visible_in_view(graph_backend: BaseGraph) -> None: + """An edge attribute written on the root must be readable through the view.""" + edge_id, _ = _graph_with_edge(graph_backend) + + view = graph_backend.filter().subgraph() + + graph_backend.update_edge_attrs(attrs={"w": 9.0}, edge_ids=[edge_id]) + + assert _edge_value_of(graph_backend, edge_id) == 9.0 + assert _edge_value_of(view, edge_id) == 9.0, "view returned a stale edge value after a root write" + + +def test_root_edge_update_all_edges_is_visible_in_view(graph_backend: BaseGraph) -> None: + """`edge_ids=None` means all edges, and must reach the view too.""" + edge_id, other_id = _graph_with_edge(graph_backend) + + view = graph_backend.filter().subgraph() + + graph_backend.update_edge_attrs(attrs={"w": 4.0}) + + assert _edge_value_of(view, edge_id) == 4.0 + assert _edge_value_of(view, other_id) == 4.0 + + +def test_root_edge_update_per_edge_values_reach_view(graph_backend: BaseGraph) -> None: + """Per-edge sequence values must land on the matching edges in the view.""" + edge_id, other_id = _graph_with_edge(graph_backend) + + view = graph_backend.filter().subgraph() + + graph_backend.update_edge_attrs(attrs={"w": [11.0, 22.0]}, edge_ids=[edge_id, other_id]) + + assert _edge_value_of(view, edge_id) == 11.0 + assert _edge_value_of(view, other_id) == 22.0 + + +def test_view_edge_write_is_applied_once(graph_backend: BaseGraph) -> None: + """Writing an edge through the view must reach both graphs exactly once.""" + edge_id, _ = _graph_with_edge(graph_backend) + + view = graph_backend.filter().subgraph() + + view.update_edge_attrs(attrs={"w": 3.0}, edge_ids=[edge_id]) + + assert _edge_value_of(graph_backend, edge_id) == 3.0 + assert _edge_value_of(view, edge_id) == 3.0 + + +def test_root_edge_update_notifies_multiple_views(graph_backend: BaseGraph) -> None: + """Every registered view holding the edge must see the new value.""" + edge_id, _ = _graph_with_edge(graph_backend) + + view_a = graph_backend.filter().subgraph() + view_b = graph_backend.filter().subgraph() + + graph_backend.update_edge_attrs(attrs={"w": 6.0}, edge_ids=[edge_id]) + + assert _edge_value_of(view_a, edge_id) == 6.0 + assert _edge_value_of(view_b, edge_id) == 6.0 + + +def test_root_edge_update_partial_overlap_keeps_positions(graph_backend: BaseGraph) -> None: + """A batch spanning in- and out-of-view edges must apply the right value. + + Per-edge values are positional, so selecting the in-view subset has to keep + the original indices rather than just the ids. + """ + graph_backend.add_node_attr_key("x", pl.Float64) + graph_backend.add_edge_attr_key("w", pl.Float64) + nodes = [graph_backend.add_node({"t": t, "x": 0.0}) for t in range(4)] + inside = graph_backend.add_edge(nodes[0], nodes[1], {"w": 0.0}) + outside = graph_backend.add_edge(nodes[2], nodes[3], {"w": 0.0}) + + view = graph_backend.filter(NodeAttr("t") <= 1).subgraph() + assert view.edge_ids() == [inside] + + graph_backend.update_edge_attrs(attrs={"w": [11.0, 22.0]}, edge_ids=[inside, outside]) + + assert _edge_value_of(view, inside) == 11.0, "view took the wrong positional value" + assert _edge_value_of(graph_backend, outside) == 22.0 + assert view._out_of_sync is False diff --git a/src/tracksdata/graph/_test/test_subgraph.py b/src/tracksdata/graph/_test/test_subgraph.py index 520bbeba..d0cc5f2a 100644 --- a/src/tracksdata/graph/_test/test_subgraph.py +++ b/src/tracksdata/graph/_test/test_subgraph.py @@ -2062,3 +2062,182 @@ def test_sql_graph_filter_borderline_node_ids(tmp_path, monkeypatch: pytest.Monk del filtered, subgraph gc.collect() assert _scratch_table_count(graph) == 0 + + +def _root_with_two_connected_nodes(graph_backend: BaseGraph) -> BaseGraph: + """A root graph with two nodes, one edge, and two attribute keys on each.""" + graph_backend.add_node_attr_key("area", default_value=0.0, dtype=pl.Float64) + graph_backend.add_node_attr_key("bar", default_value=0.0, dtype=pl.Float64) + graph_backend.add_edge_attr_key("weight", default_value=0.0, dtype=pl.Float64) + graph_backend.add_edge_attr_key("cost", default_value=0.0, dtype=pl.Float64) + source = graph_backend.add_node({"t": 0, "area": 1.0, "bar": 1.0}) + target = graph_backend.add_node({"t": 1, "area": 2.0, "bar": 2.0}) + graph_backend.add_edge(source, target, {"weight": 1.0, "cost": 1.0}) + return graph_backend + + +def test_update_root_node_key_outside_view_attr_keys(graph_backend: BaseGraph) -> None: + """A view tracking a subset of keys tolerates root updates to the keys it excluded. + + A view built with an explicit `node_attr_keys` has no local column for the + keys it left out, so the update must not be propagated into it. + """ + root = _root_with_two_connected_nodes(graph_backend) + view = root.filter().subgraph(node_attr_keys=["area"]) + + assert "bar" not in view.node_attr_keys() + + root.update_node_attrs(attrs={"bar": [9.0]}, node_ids=[root.node_ids()[0]]) + + assert root.node_attrs(attr_keys=["bar"])["bar"].to_list() == [9.0, 2.0] + assert "bar" not in view.node_attr_keys() + # the keys the view does track are still maintained + root.update_node_attrs(attrs={"area": [7.0]}, node_ids=[root.node_ids()[0]]) + assert view.node_attrs(attr_keys=["area"])["area"].to_list() == [7.0, 2.0] + + +def test_update_root_edge_key_outside_view_attr_keys(graph_backend: BaseGraph) -> None: + """The edge counterpart of `test_update_root_node_key_outside_view_attr_keys`.""" + root = _root_with_two_connected_nodes(graph_backend) + view = root.filter().subgraph(edge_attr_keys=["weight"]) + + assert "cost" not in view.edge_attr_keys() + + root.update_edge_attrs(attrs={"cost": [9.0]}, edge_ids=[root.edge_ids()[0]]) + + assert root.edge_attrs(attr_keys=["cost"])["cost"].to_list() == [9.0] + assert "cost" not in view.edge_attr_keys() + # the keys the view does track are still maintained + root.update_edge_attrs(attrs={"weight": [7.0]}, edge_ids=[root.edge_ids()[0]]) + assert view.edge_attrs(attr_keys=["weight"])["weight"].to_list() == [7.0] + + +def test_add_node_attr_key_on_root_reaches_live_view(graph_backend: BaseGraph) -> None: + """A key registered on the root must reach the views already derived from it. + + A rustworkx-rooted view reports the root's keys and shares its attribute + dicts, so it picks the key up for free. A SQLGraph-rooted view holds its own + copy of both and has to be told. + """ + root = _root_with_two_connected_nodes(graph_backend) + view = root.filter().subgraph() + + root.add_node_attr_key("foo", default_value=-1, dtype=pl.Int64) + + assert "foo" in root.node_attr_keys() + assert "foo" in view.node_attr_keys() + assert view.node_attrs(attr_keys=["foo"])["foo"].to_list() == [-1, -1] + + # the view's local store accepts writes to the new key, on either side + root.update_node_attrs(attrs={"foo": [7]}, node_ids=[root.node_ids()[0]]) + assert view.node_attrs(attr_keys=["foo"])["foo"].to_list() == [7, -1] + + +def test_add_edge_attr_key_on_root_reaches_live_view(graph_backend: BaseGraph) -> None: + """The edge counterpart of `test_add_node_attr_key_on_root_reaches_live_view`.""" + root = _root_with_two_connected_nodes(graph_backend) + view = root.filter().subgraph() + + root.add_edge_attr_key("w", default_value=-1.0, dtype=pl.Float64) + + assert "w" in root.edge_attr_keys() + assert "w" in view.edge_attr_keys() + assert view.edge_attrs(attr_keys=["w"])["w"].to_list() == [-1.0] + + root.update_edge_attrs(attrs={"w": [1.5]}, edge_ids=[root.edge_ids()[0]]) + assert view.edge_attrs(attr_keys=["w"])["w"].to_list() == [1.5] + + +def test_add_attr_key_on_view_reaches_sibling_view(graph_backend: BaseGraph) -> None: + """Registering through one view must reach the other views of the same root.""" + root = _root_with_two_connected_nodes(graph_backend) + view_a = root.filter().subgraph() + view_b = root.filter().subgraph() + + view_a.add_node_attr_key("foo", default_value=-1, dtype=pl.Int64) + view_a.add_edge_attr_key("w", default_value=-1.0, dtype=pl.Float64) + + assert "foo" in view_b.node_attr_keys() + assert "w" in view_b.edge_attr_keys() + assert view_b.node_attrs(attr_keys=["foo"])["foo"].to_list() == [-1, -1] + assert view_b.edge_attrs(attr_keys=["w"])["w"].to_list() == [-1.0] + + +def test_remove_node_attr_key_on_root_reaches_live_view(graph_backend: BaseGraph) -> None: + """The remove counterpart of `test_add_node_attr_key_on_root_reaches_live_view`. + + Dropping a key on the root must drop it from the views already derived from + it, both from the reported schema and from the view's local attribute store. + """ + root = _root_with_two_connected_nodes(graph_backend) + view = root.filter().subgraph() + + root.remove_node_attr_key("bar") + + assert "bar" not in root.node_attr_keys() + assert "bar" not in view.node_attr_keys() + with pytest.raises(KeyError): + view.node_attrs(attr_keys=["bar"]) + # the keys that remain are untouched + assert view.node_attrs(attr_keys=["area"])["area"].to_list() == [1.0, 2.0] + + +def test_remove_edge_attr_key_on_root_reaches_live_view(graph_backend: BaseGraph) -> None: + """The edge counterpart of `test_remove_node_attr_key_on_root_reaches_live_view`.""" + root = _root_with_two_connected_nodes(graph_backend) + view = root.filter().subgraph() + + root.remove_edge_attr_key("cost") + + assert "cost" not in root.edge_attr_keys() + assert "cost" not in view.edge_attr_keys() + with pytest.raises(KeyError): + view.edge_attrs(attr_keys=["cost"]) + assert view.edge_attrs(attr_keys=["weight"])["weight"].to_list() == [1.0] + + +def test_remove_node_attr_key_on_root_updates_pinned_view_keys(graph_backend: BaseGraph) -> None: + """A view pinning an explicit key list must not keep reporting a removed key. + + The pinned list is the view's own copy of the schema, so a root removal has + to be applied to it -- otherwise the view advertises a column that no longer + exists anywhere. + """ + root = _root_with_two_connected_nodes(graph_backend) + view = root.filter().subgraph(node_attr_keys=["area", "bar"], edge_attr_keys=["weight", "cost"]) + + assert "bar" in view.node_attr_keys() + assert "cost" in view.edge_attr_keys() + + root.remove_node_attr_key("bar") + root.remove_edge_attr_key("cost") + + assert "bar" not in view.node_attr_keys() + assert "cost" not in view.edge_attr_keys() + # node_attrs() with no attr_keys uses the pinned list, so a stale entry there + # surfaces as a failure to materialize the view at all + assert "bar" not in view.node_attrs().columns + assert "cost" not in view.edge_attrs().columns + + +def test_remove_attr_key_on_view_reaches_sibling_view(graph_backend: BaseGraph) -> None: + """The remove counterpart of `test_add_attr_key_on_view_reaches_sibling_view`. + + Removing through one view propagates up to the root, so it must also reach + the root's other views. + """ + root = _root_with_two_connected_nodes(graph_backend) + view_a = root.filter().subgraph() + view_b = root.filter().subgraph() + + view_a.remove_node_attr_key("bar") + view_a.remove_edge_attr_key("cost") + + assert "bar" not in root.node_attr_keys() + assert "cost" not in root.edge_attr_keys() + assert "bar" not in view_b.node_attr_keys() + assert "cost" not in view_b.edge_attr_keys() + with pytest.raises(KeyError): + view_b.node_attrs(attr_keys=["bar"]) + with pytest.raises(KeyError): + view_b.edge_attrs(attr_keys=["cost"])