diff --git a/src/tracksdata/graph/_base_graph.py b/src/tracksdata/graph/_base_graph.py index d7fbfdba..c76d0fe8 100644 --- a/src/tracksdata/graph/_base_graph.py +++ b/src/tracksdata/graph/_base_graph.py @@ -197,6 +197,22 @@ def _maintain_views_edge_attrs( 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) + @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 aed9ed79..1a885f63 100644 --- a/src/tracksdata/graph/_graph_view.py +++ b/src/tracksdata/graph/_graph_view.py @@ -309,30 +309,11 @@ 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: self._root.remove_node_attr_key(key) if self._node_attr_keys is not None and key in self._node_attr_keys: @@ -351,28 +332,47 @@ 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) + 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. - # 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 + 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: self._root.remove_edge_attr_key(key) diff --git a/src/tracksdata/graph/_rustworkx_graph.py b/src/tracksdata/graph/_rustworkx_graph.py index 92451e90..500ef23e 100644 --- a/src/tracksdata/graph/_rustworkx_graph.py +++ b/src/tracksdata/graph/_rustworkx_graph.py @@ -1076,6 +1076,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. @@ -1116,6 +1118,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. diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index ef54229c..56c31473 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/src/tracksdata/graph/_sql_graph.py @@ -1966,6 +1966,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") @@ -1993,6 +1995,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") diff --git a/src/tracksdata/graph/_test/test_subgraph.py b/src/tracksdata/graph/_test/test_subgraph.py index be88412e..9a8369b9 100644 --- a/src/tracksdata/graph/_test/test_subgraph.py +++ b/src/tracksdata/graph/_test/test_subgraph.py @@ -2093,3 +2093,54 @@ def test_update_root_edge_key_outside_view_attr_keys(graph_backend: BaseGraph) - # 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]