From fa423935f98c74c4ba02dfeb74b402032fa4c548 Mon Sep 17 00:00:00 2001 From: Teun Huijben Date: Tue, 15 Sep 2026 15:30:51 -0700 Subject: [PATCH 1/6] let graphview read missing attributes from root --- src/tracksdata/graph/_graph_view.py | 119 +++++++- src/tracksdata/graph/_rustworkx_graph.py | 2 + src/tracksdata/graph/_sql_graph.py | 2 + .../graph/_test/test_root_fallback.py | 288 ++++++++++++++++++ src/tracksdata/graph/filters/_base_filter.py | 6 + .../graph/filters/_indexed_filter.py | 2 + 6 files changed, 412 insertions(+), 7 deletions(-) create mode 100644 src/tracksdata/graph/_test/test_root_fallback.py diff --git a/src/tracksdata/graph/_graph_view.py b/src/tracksdata/graph/_graph_view.py index dfe7d8c8..84f8970f 100644 --- a/src/tracksdata/graph/_graph_view.py +++ b/src/tracksdata/graph/_graph_view.py @@ -13,6 +13,7 @@ from tracksdata.graph._mapped_graph_mixin import MappedGraphMixin from tracksdata.graph._rustworkx_graph import IndexedRXGraph, RustWorkXGraph, RXFilter from tracksdata.graph.filters._indexed_filter import IndexRXFilter +from tracksdata.utils._dataframe import unpack_array_attrs from tracksdata.utils._dtypes import AttrSchema from tracksdata.utils._signal import ( emit_node_added_events, @@ -107,6 +108,13 @@ class GraphView(MappedGraphMixin, RustWorkXGraph): sync : bool, default True Whether to automatically synchronize changes in the view. By default only the root graph is updated. + root_fallback : bool, default False + When True, reading a node attribute key this view does not hold is + served from the root graph instead of raising. This is what makes a + view built with a reduced `node_attr_keys` usable for the keys it left + out -- at the cost of going to the root's storage for them, which for a + SQL root is a query per read. Read path only -- writes already go + straight to the root for keys the view does not hold. Attributes ---------- @@ -157,6 +165,7 @@ def __init__( mode: ViewMode = ViewMode.WRITE_THROUGH, node_attr_keys: list[str] | None = None, edge_attr_keys: list[str] | None = None, + root_fallback: bool = False, ) -> None: # Initialize RustWorkXGraph RustWorkXGraph.__init__(self, rx_graph=None) # rx_graph is not used to avoid initialization @@ -211,6 +220,10 @@ def __init__( ] self._edge_attr_keys = list(dict.fromkeys(self._edge_attr_keys)) + # When set, a read for a node attribute key this view does not hold is + # served from the root instead of raising. See `_node_attrs_from_node_ids`. + self._root_fallback = root_fallback + # use parent graph overlaps self._overlaps = None @@ -363,6 +376,35 @@ def node_attr_keys(self, return_ids: bool = False) -> list[str]: pass return keys + def _validate_attr_keys( + self, + attr_keys: Sequence[str] | str | None, + mode: Literal["node", "edge"], + ) -> None: + """ + Same validation as any other graph, plus a pointer when the key is only + missing because this view was built without it. + + A view built with an explicit `node_attr_keys` list is otherwise a trap: + the key exists, the root has it, and the error says only that this graph + does not. + """ + try: + super()._validate_attr_keys(attr_keys, mode) + except KeyError as err: + if mode != "node" or self._root_fallback: + raise + if isinstance(attr_keys, str): + attr_keys = [attr_keys] + missing = set(attr_keys) - set(self.node_attr_keys(return_ids=True)) + on_root = sorted(missing & set(self._root.node_attr_keys(return_ids=True))) + if not on_root: + raise + raise KeyError( + f"{err.args[0]} The root graph does hold {on_root}; build the view with " + "`root_fallback=True` to read those through it." + ) from None + def edge_attr_keys(self, return_ids: bool = False) -> list[str]: """ Get the keys of the attributes of the edges. @@ -1037,6 +1079,26 @@ def predecessors( raise RuntimeError("Out of sync graph view cannot be used to get predecessors") return super().predecessors(node_ids, attr_keys, return_attrs=return_attrs) + def _split_node_attr_keys(self, attr_keys: Sequence[str]) -> tuple[list[str], list[str]]: + """ + Partition requested node attribute keys into those this view holds and + those only the root holds. + + Only ever returns root keys when the view was built with + ``root_fallback=True``; otherwise every key is reported as local and the + usual validation in the local read path rejects the ones that are not. + Keys that exist on neither the view nor the root still raise, here via + the root's own validation. + """ + if not self._root_fallback: + return list(attr_keys), [] + + local = set(self.node_attr_keys(return_ids=True)) + root_keys = [k for k in attr_keys if k not in local] + if root_keys: + self._root._validate_attr_keys(root_keys, "node") + return [k for k in attr_keys if k in local], root_keys + def _node_attrs_from_node_ids( self, *, @@ -1044,14 +1106,57 @@ def _node_attrs_from_node_ids( attr_keys: Sequence[str] | str | None = None, unpack: bool = False, ) -> pl.DataFrame: - node_dfs = super()._node_attrs_from_node_ids( - node_ids=self._map_to_local(node_ids), - attr_keys=attr_keys, - unpack=unpack, - ) - node_dfs = self._map_df_to_external( - node_dfs, [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET] + id_keys = [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET] + + if attr_keys is None: + # "every key" means every key this view holds -- a fallback would + # silently pull in the columns the view was built to leave out. + node_dfs = super()._node_attrs_from_node_ids( + node_ids=self._map_to_local(node_ids), + attr_keys=None, + unpack=unpack, + ) + return self._map_df_to_external(node_dfs, id_keys) + + if isinstance(attr_keys, str): + attr_keys = [attr_keys] + attr_keys = list(dict.fromkeys(attr_keys)) + + local_keys, root_keys = self._split_node_attr_keys(attr_keys) + + if not root_keys: + node_dfs = super()._node_attrs_from_node_ids( + node_ids=self._map_to_local(node_ids), + attr_keys=attr_keys, + unpack=unpack, + ) + return self._map_df_to_external(node_dfs, id_keys) + + node_id_key = DEFAULT_ATTR_KEYS.NODE_ID + external_ids = self.node_ids() if node_ids is None else list(node_ids) + + # NODE_ID is needed as the join key even when the caller did not ask for + # it, and `unpack` is deferred until after the join so that it sees the + # same columns it would have on a view holding all of them. + local_df = super()._node_attrs_from_node_ids( + node_ids=self._map_to_local(external_ids), + attr_keys=[node_id_key, *local_keys], + unpack=False, ) + local_df = self._map_df_to_external(local_df, id_keys) + + # Ask the root for exactly the rows this view covers -- never for all of + # its rows, which include the ones this view deliberately dropped. + root_df = self._root.filter(node_ids=external_ids).node_attrs(attr_keys=[node_id_key, *root_keys]) + + # `maintain_order="left"` is not optional: callers zip columns + # positionally, so a reordered join would pair the wrong values. + node_dfs = local_df.join(root_df, on=node_id_key, how="left", maintain_order="left") + node_dfs = node_dfs.select(attr_keys) + + if unpack: + node_dfs = unpack_array_attrs(node_dfs) + return node_dfs def node_attrs( diff --git a/src/tracksdata/graph/_rustworkx_graph.py b/src/tracksdata/graph/_rustworkx_graph.py index 98247e7c..9df10c0a 100644 --- a/src/tracksdata/graph/_rustworkx_graph.py +++ b/src/tracksdata/graph/_rustworkx_graph.py @@ -368,6 +368,7 @@ def subgraph( edge_attr_keys: Sequence[str] | None = None, *, mode: "ViewMode | None" = None, + root_fallback: bool = False, ) -> "GraphView": from tracksdata.graph._graph_view import GraphView, ViewMode @@ -387,6 +388,7 @@ def subgraph( mode=mode if mode is not None else ViewMode.WRITE_THROUGH, node_attr_keys=node_attr_keys, edge_attr_keys=edge_attr_keys, + root_fallback=root_fallback, ) return graph_view diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index d4ab8663..433dad73 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/src/tracksdata/graph/_sql_graph.py @@ -524,6 +524,7 @@ def subgraph( edge_attr_keys: Sequence[str] | None = None, *, mode: "ViewMode | None" = None, + root_fallback: bool = False, ) -> "GraphView": from tracksdata.graph._graph_view import GraphView, ViewMode @@ -582,6 +583,7 @@ def subgraph( mode=mode if mode is not None else ViewMode.WRITE_THROUGH, node_attr_keys=node_attr_keys, edge_attr_keys=edge_attr_keys, + root_fallback=root_fallback, ) return graph diff --git a/src/tracksdata/graph/_test/test_root_fallback.py b/src/tracksdata/graph/_test/test_root_fallback.py new file mode 100644 index 00000000..f8d96010 --- /dev/null +++ b/src/tracksdata/graph/_test/test_root_fallback.py @@ -0,0 +1,288 @@ +"""Tests for `GraphView(root_fallback=True)`: reading node attribute keys a +partial view was built without, by falling back to its root graph.""" + +import numpy as np +import polars as pl +import pytest + +from tracksdata.array import GraphArrayView +from tracksdata.attrs import NodeAttr +from tracksdata.constants import DEFAULT_ATTR_KEYS +from tracksdata.graph import BaseGraph, GraphView +from tracksdata.nodes._mask import Mask + +LEAN_KEYS = ["x", "y"] + + +def _populate(graph: BaseGraph) -> BaseGraph: + """A graph with a cheap key set and one key a lean view will leave out.""" + graph.add_node_attr_key("x", dtype=pl.Float64) + graph.add_node_attr_key("y", dtype=pl.Float64) + graph.add_node_attr_key("label", dtype=pl.String, default_value="") + graph.add_node_attr_key("solution", dtype=pl.Boolean, default_value=False) + graph.add_node_attr_key("vec", dtype=pl.Array(pl.Int64, 2)) + + for i in range(6): + graph.add_node( + { + DEFAULT_ATTR_KEYS.T: i, + "x": float(i), + "y": float(i * 10), + "label": f"n{i}", + "solution": i % 2 == 0, + "vec": np.array([i, -i]), + } + ) + return graph + + +def _views(graph: BaseGraph, **kwargs) -> tuple[GraphView, GraphView]: + """A lean view and the equivalent full view, over the same nodes.""" + lean = graph.filter(NodeAttr("solution") == True).subgraph(node_attr_keys=LEAN_KEYS, root_fallback=True, **kwargs) + full = graph.filter(NodeAttr("solution") == True).subgraph() + return lean, full + + +# --- correctness of the fallback ------------------------------------------ + + +def test_fallback_matches_full_view(graph_backend: BaseGraph) -> None: + lean, full = _views(_populate(graph_backend)) + + assert "label" not in lean.node_attr_keys() + + keys = [DEFAULT_ATTR_KEYS.NODE_ID, "label"] + assert lean.node_attrs(attr_keys=keys).equals(full.node_attrs(attr_keys=keys)) + + +def test_mixed_request_keeps_requested_order(graph_backend: BaseGraph) -> None: + lean, full = _views(_populate(graph_backend)) + + keys = ["label", "x", DEFAULT_ATTR_KEYS.NODE_ID, "y"] + df = lean.node_attrs(attr_keys=keys) + + assert df.columns == keys + assert df.equals(full.node_attrs(attr_keys=keys)) + + +def test_fallback_without_node_id_requested(graph_backend: BaseGraph) -> None: + """NODE_ID is the join key but must not leak into the result.""" + lean, full = _views(_populate(graph_backend)) + + df = lean.node_attrs(attr_keys=["label", "x"]) + + assert df.columns == ["label", "x"] + assert df.equals(full.node_attrs(attr_keys=["label", "x"])) + + +def test_fallback_single_str_attr_key(graph_backend: BaseGraph) -> None: + lean, full = _views(_populate(graph_backend)) + + df = lean.node_attrs(attr_keys="label") + + assert df.columns == ["label"] + assert df.equals(full.node_attrs(attr_keys="label")) + + +def test_fallback_unpack(graph_backend: BaseGraph) -> None: + lean, full = _views(_populate(graph_backend)) + + keys = ["x", "vec"] + df = lean.node_attrs(attr_keys=keys, unpack=True) + + assert df.equals(full.node_attrs(attr_keys=keys, unpack=True)) + assert "vec" not in df.columns + + +def test_fallback_node_ids_subset(graph_backend: BaseGraph) -> None: + lean, full = _views(_populate(graph_backend)) + subset = lean.node_ids()[:2] + + keys = [DEFAULT_ATTR_KEYS.NODE_ID, "label"] + df = lean.filter(node_ids=subset).node_attrs(attr_keys=keys) + + assert df[DEFAULT_ATTR_KEYS.NODE_ID].to_list() == subset + assert df.equals(full.filter(node_ids=subset).node_attrs(attr_keys=keys)) + + +def test_fallback_never_returns_rows_outside_the_view(graph_backend: BaseGraph) -> None: + """The root holds the `solution=False` nodes the view dropped.""" + graph = _populate(graph_backend) + lean, _ = _views(graph) + + assert graph.num_nodes() > lean.num_nodes() + + df = lean.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, "label"]) + + assert df.height == lean.num_nodes() + assert set(df[DEFAULT_ATTR_KEYS.NODE_ID].to_list()) == set(lean.node_ids()) + + +def test_fallback_row_order_matches_local_columns(graph_backend: BaseGraph) -> None: + """Callers zip columns positionally; the join must not reorder rows.""" + lean, _ = _views(_populate(graph_backend)) + + df = lean.node_attrs(attr_keys=["x", "label"]) + + for x, label in zip(df["x"], df["label"], strict=True): + assert label == f"n{int(x)}" + + +def test_fallback_after_local_update_is_not_stale(graph_backend: BaseGraph) -> None: + lean, _ = _views(_populate(graph_backend)) + node_ids = lean.node_ids() + + lean.update_node_attrs(node_ids=node_ids, attrs={"x": 99.0}) + + df = lean.node_attrs(attr_keys=["x", "label"]) + assert df["x"].to_list() == [99.0] * len(node_ids) + assert df["label"].to_list() == sorted(df["label"].to_list()) + + +def test_fallback_skips_node_removed_from_view(graph_backend: BaseGraph) -> None: + graph = _populate(graph_backend) + lean, _ = _views(graph) + + removed = lean.node_ids()[0] + lean.remove_node_from_view(removed) + + df = lean.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, "label"]) + + assert removed not in df[DEFAULT_ATTR_KEYS.NODE_ID].to_list() + assert graph.has_node(removed) + + +def test_attr_keys_none_stays_lean(graph_backend: BaseGraph) -> None: + """`None` means "what this view holds", not "everything the root has".""" + lean, _ = _views(_populate(graph_backend)) + + assert "label" not in lean.node_attrs().columns + + +# --- error behaviour preserved -------------------------------------------- + + +def test_unknown_key_still_raises(graph_backend: BaseGraph) -> None: + lean, _ = _views(_populate(graph_backend)) + + with pytest.raises(KeyError, match="not found"): + lean.node_attrs(attr_keys=["nowhere"]) + + +def test_without_fallback_missing_key_raises_with_hint(graph_backend: BaseGraph) -> None: + graph = _populate(graph_backend) + view = graph.filter(NodeAttr("solution") == True).subgraph(node_attr_keys=LEAN_KEYS) + + with pytest.raises(KeyError, match="root_fallback=True"): + view.node_attrs(attr_keys=["label"]) + + +def test_without_fallback_unknown_key_has_no_hint(graph_backend: BaseGraph) -> None: + graph = _populate(graph_backend) + view = graph.filter(NodeAttr("solution") == True).subgraph(node_attr_keys=LEAN_KEYS) + + with pytest.raises(KeyError) as excinfo: + view.node_attrs(attr_keys=["nowhere"]) + + assert "root_fallback" not in str(excinfo.value) + + +def test_update_on_missing_key_writes_to_root_and_reads_back(graph_backend: BaseGraph) -> None: + """A write to a non-local key already goes to the root (pre-existing + behaviour: `_update_local_node_attrs` skips keys the view does not hold). + Fallback cannot make that stale -- there is no local copy to disagree.""" + graph = _populate(graph_backend) + lean, _ = _views(graph) + node_ids = lean.node_ids() + + lean.update_node_attrs(node_ids=node_ids, attrs={"label": "written"}) + + assert lean.node_attrs(attr_keys=["label"])["label"].to_list() == ["written"] * len(node_ids) + assert graph.filter(node_ids=node_ids).node_attrs(attr_keys=["label"])["label"].to_list() == ["written"] * len( + node_ids + ) + + +def test_node_attr_keys_reports_only_local_keys(graph_backend: BaseGraph) -> None: + lean, _ = _views(_populate(graph_backend)) + + assert sorted(lean.node_attr_keys()) == sorted([*LEAN_KEYS, DEFAULT_ATTR_KEYS.T]) + + +def test_fallback_is_off_by_default(graph_backend: BaseGraph) -> None: + view = _populate(graph_backend).filter(NodeAttr("solution") == True).subgraph(node_attr_keys=LEAN_KEYS) + + assert view._root_fallback is False + + +# --- end to end through GraphArrayView ------------------------------------- + + +def test_graph_array_view_on_lean_view(graph_backend: BaseGraph) -> None: + """The case the feature exists for: render masks the view does not hold.""" + graph = graph_backend + graph.add_node_attr_key("label", dtype=pl.Int64) + graph.add_node_attr_key(DEFAULT_ATTR_KEYS.MASK, pl.Object) + graph.add_node_attr_key(DEFAULT_ATTR_KEYS.BBOX, pl.Array(pl.Int64, 4)) + graph.add_node_attr_key("solution", dtype=pl.Boolean, default_value=False) + + for i, solution in enumerate([True, False, True]): + mask = Mask( + np.ones((2, 2), dtype=bool), + bbox=np.array([10 * i, 20, 10 * i + 2, 22]), + ) + graph.add_node( + { + DEFAULT_ATTR_KEYS.T: 0, + "label": i + 1, + DEFAULT_ATTR_KEYS.MASK: mask, + DEFAULT_ATTR_KEYS.BBOX: mask.bbox, + "solution": solution, + } + ) + + node_filter = graph.filter(NodeAttr("solution") == True) + # bbox stays local: the view's own R-tree is built from it. + lean = node_filter.subgraph(node_attr_keys=["label", DEFAULT_ATTR_KEYS.BBOX], root_fallback=True) + full = node_filter.subgraph() + + shape = (1, 100, 100) + lean_array = np.asarray(GraphArrayView(graph=lean, shape=shape, attr_key="label")[0]) + full_array = np.asarray(GraphArrayView(graph=full, shape=shape, attr_key="label")[0]) + + assert lean_array.max() > 0 + np.testing.assert_array_equal(lean_array, full_array) + + +# --- regression guard: the excluded column is never materialized ------------ + + +class _ExplodingBlob: + """Unpickles only if something actually reads the column it lives in.""" + + def __setstate__(self, state: dict) -> None: + raise AssertionError("excluded blob column was read while building the lean view") + + def __getstate__(self) -> dict: + return {} + + +def test_lean_view_does_not_read_the_excluded_column() -> None: + from tracksdata.graph import SQLGraph + + graph = SQLGraph(drivername="sqlite", database=":memory:") + graph.add_node_attr_key("x", dtype=pl.Float64) + graph.add_node_attr_key("blob", dtype=pl.Object) + + for i in range(4): + graph.add_node({DEFAULT_ATTR_KEYS.T: i, "x": float(i), "blob": _ExplodingBlob()}) + + # Would raise if `blob` were fetched and unpickled here. + lean = graph.filter().subgraph(node_attr_keys=["x"], root_fallback=True) + + assert "blob" not in lean.node_attr_keys() + # ... and it is genuinely absent from the view's own storage, not just hidden. + assert all("blob" not in payload for payload in lean.rx_graph.nodes()) + + with pytest.raises(AssertionError, match="excluded blob column"): + lean.node_attrs(attr_keys=["blob"]) diff --git a/src/tracksdata/graph/filters/_base_filter.py b/src/tracksdata/graph/filters/_base_filter.py index e6ea6b77..acafb4c2 100644 --- a/src/tracksdata/graph/filters/_base_filter.py +++ b/src/tracksdata/graph/filters/_base_filter.py @@ -51,6 +51,7 @@ def subgraph( edge_attr_keys: list[str] | None = None, *, mode: "ViewMode | None" = None, + root_fallback: bool = False, ) -> "GraphView": """ Get a subgraph of the graph resulting from the filter. @@ -59,6 +60,11 @@ def subgraph( How the resulting view relates to its root graph (write-through only, or live-updating). Defaults to ``ViewMode.WRITE_THROUGH`` when None. + root_fallback : bool + When True, reading a node attribute key the resulting view does not + hold is served from the root graph instead of raising. Off by + default: the fetch goes to the root's storage, which for a SQL root + means a query per read. """ @abc.abstractmethod diff --git a/src/tracksdata/graph/filters/_indexed_filter.py b/src/tracksdata/graph/filters/_indexed_filter.py index 8efea0d9..be06181e 100644 --- a/src/tracksdata/graph/filters/_indexed_filter.py +++ b/src/tracksdata/graph/filters/_indexed_filter.py @@ -54,6 +54,7 @@ def subgraph( edge_attr_keys: Sequence[str] | str | None = None, *, mode: "ViewMode | None" = None, + root_fallback: bool = False, ) -> "GraphView": from tracksdata.graph._graph_view import GraphView, ViewMode @@ -77,6 +78,7 @@ def subgraph( mode=mode if mode is not None else ViewMode.WRITE_THROUGH, node_attr_keys=node_attr_keys, edge_attr_keys=edge_attr_keys, + root_fallback=root_fallback, ) return graph_view From e2513cc9509969878118a5b2cc5e5922cb6265af Mon Sep 17 00:00:00 2001 From: Teun Huijben Date: Wed, 16 Sep 2026 11:54:15 -0700 Subject: [PATCH 2/6] implement logic for nested views --- .../graph/_test/test_root_fallback.py | 65 +++++++++++++++++++ .../graph/filters/_indexed_filter.py | 32 +++++++++ 2 files changed, 97 insertions(+) diff --git a/src/tracksdata/graph/_test/test_root_fallback.py b/src/tracksdata/graph/_test/test_root_fallback.py index f8d96010..feef566a 100644 --- a/src/tracksdata/graph/_test/test_root_fallback.py +++ b/src/tracksdata/graph/_test/test_root_fallback.py @@ -286,3 +286,68 @@ def test_lean_view_does_not_read_the_excluded_column() -> None: with pytest.raises(AssertionError, match="excluded blob column"): lean.node_attrs(attr_keys=["blob"]) + + +# --- a view of a lean view --------------------------------------------------- + + +def test_nested_subgraph_inherits_the_lean_key_set(graph_backend: BaseGraph) -> None: + """A child's rx payloads are copied from its parent, so it holds no more + than the parent did -- and must not advertise otherwise.""" + lean, _ = _views(_populate(graph_backend)) + + nested = lean.filter(node_ids=lean.node_ids()[:2]).subgraph() + + assert "label" not in nested.node_attr_keys() + assert nested.node_attrs(attr_keys=["label"])["label"].to_list() == ["n0", "n2"] + + +def test_nested_subgraph_without_fallback_raises(graph_backend: BaseGraph) -> None: + """The silent-defaults case: it must fail loudly, not return schema defaults.""" + graph = _populate(graph_backend) + lean = graph.filter(NodeAttr("solution") == True).subgraph(node_attr_keys=LEAN_KEYS) + + nested = lean.filter(node_ids=lean.node_ids()[:2]).subgraph() + + with pytest.raises(KeyError, match="root_fallback=True"): + nested.node_attrs(attr_keys=["label"]) + + +def test_nested_subgraph_explicit_key_the_parent_lacks(graph_backend: BaseGraph) -> None: + """Naming an excluded key explicitly must not re-claim it locally.""" + lean, _ = _views(_populate(graph_backend)) + + nested = lean.filter(node_ids=lean.node_ids()[:2]).subgraph(node_attr_keys=["label"]) + + assert "label" not in nested.node_attr_keys() + assert nested.node_attrs(attr_keys=["label"])["label"].to_list() == ["n0", "n2"] + + +def test_nested_subgraph_of_full_view_is_unclamped(graph_backend: BaseGraph) -> None: + """A parent holding everything clamps nothing.""" + graph = _populate(graph_backend) + full = graph.filter(NodeAttr("solution") == True).subgraph() + + nested = full.filter(node_ids=full.node_ids()[:2]).subgraph(node_attr_keys=["label"]) + + assert "label" in nested.node_attr_keys() + assert nested.node_attrs(attr_keys=["label"])["label"].to_list() == ["n0", "n2"] + + +def test_nested_subgraph_key_on_neither_parent_nor_root(graph_backend: BaseGraph) -> None: + """A key that exists nowhere raises plainly -- pointing at `root_fallback` + would suggest a flag that cannot help.""" + graph = _populate(graph_backend) + + for fallback in (True, False): + lean = graph.filter(NodeAttr("solution") == True).subgraph( + node_attr_keys=LEAN_KEYS, + root_fallback=fallback, + ) + nested = lean.filter(node_ids=lean.node_ids()[:2]).subgraph() + + with pytest.raises(KeyError) as excinfo: + nested.node_attrs(attr_keys=["nowhere"]) + + assert "nowhere" in str(excinfo.value) + assert "root_fallback" not in str(excinfo.value) diff --git a/src/tracksdata/graph/filters/_indexed_filter.py b/src/tracksdata/graph/filters/_indexed_filter.py index be06181e..521b5cfe 100644 --- a/src/tracksdata/graph/filters/_indexed_filter.py +++ b/src/tracksdata/graph/filters/_indexed_filter.py @@ -16,6 +16,30 @@ from tracksdata.graph._graph_view import GraphView, ViewMode +def _clamp_to_parent( + requested: Sequence[str] | str | None, + parent_keys: list[str] | None, +) -> Sequence[str] | str | None: + """ + Limit the keys a child view claims to hold to the ones its parent held. + + ``parent_keys is None`` means the parent held everything its own root has, + so there is nothing to clamp. Otherwise requested keys the parent did not + hold are dropped rather than claimed: the child's rx payloads are copied + from the parent, so a key the parent lacked has no values here either. + Claiming it makes reads return schema defaults instead of raising or + falling back to the root. + """ + if parent_keys is None: + return requested + if requested is None: + return parent_keys + if isinstance(requested, str): + requested = [requested] + held = set(parent_keys) + return [k for k in requested if k in held] + + class IndexRXFilter(RXFilter): _graph: "GraphView | IndexedRXGraph" @@ -69,7 +93,15 @@ def subgraph( root = self._graph if hasattr(self._graph, "_root"): + # A view of a partial view is flattened onto the shared root, but it + # holds only what its parent held -- the rx payloads above were + # copied from the parent, not re-read from the root. Both the key + # lists and the fallback flag have to be inherited, or the child + # advertises the root's columns while holding no values for them. root = self._graph._root + node_attr_keys = _clamp_to_parent(node_attr_keys, self._graph._node_attr_keys) + edge_attr_keys = _clamp_to_parent(edge_attr_keys, self._graph._edge_attr_keys) + root_fallback = root_fallback or self._graph._root_fallback graph_view = GraphView( rx_graph, From 47fbedd657fe6737b1148cc9b9ef9977a73f8ddc Mon Sep 17 00:00:00 2001 From: Caroline Malin-Mayor Date: Sat, 17 Oct 2026 09:18:29 -0400 Subject: [PATCH 3/6] Add override and docstring --- src/tracksdata/graph/_graph_view.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/tracksdata/graph/_graph_view.py b/src/tracksdata/graph/_graph_view.py index 84f8970f..be3eda2f 100644 --- a/src/tracksdata/graph/_graph_view.py +++ b/src/tracksdata/graph/_graph_view.py @@ -6,6 +6,7 @@ import numpy as np import polars as pl import rustworkx as rx +from typing_extensions import override from tracksdata.attrs import Filter from tracksdata.constants import DEFAULT_ATTR_KEYS @@ -376,6 +377,7 @@ def node_attr_keys(self, return_ids: bool = False) -> list[str]: pass return keys + @override def _validate_attr_keys( self, attr_keys: Sequence[str] | str | None, @@ -388,6 +390,22 @@ def _validate_attr_keys( A view built with an explicit `node_attr_keys` list is otherwise a trap: the key exists, the root has it, and the error says only that this graph does not. + + Parameters + ---------- + attr_keys : Sequence[str] | str | None + The attribute keys to validate. `None` means "all keys" and is always valid. + mode : Literal["node", "edge"] + Whether to validate against node or edge attribute keys. The root-fallback + hint only ever applies to `"node"`; edge keys raise exactly as the base + implementation does. + + Raises + ------ + KeyError + If any key is not a declared attribute key of this view. When the key is + held by the root but excluded from this view, and `root_fallback` was not + set, the message also names `root_fallback=True` as the fix. """ try: super()._validate_attr_keys(attr_keys, mode) From f8ab4edf09119defacb25298e7ed08cdf6625154 Mon Sep 17 00:00:00 2001 From: Caroline Malin-Mayor Date: Sat, 17 Oct 2026 09:34:40 -0400 Subject: [PATCH 4/6] Split node id to attrs logic into heleprs --- src/tracksdata/graph/_graph_view.py | 161 +++++++++++++++++++++++----- 1 file changed, 136 insertions(+), 25 deletions(-) diff --git a/src/tracksdata/graph/_graph_view.py b/src/tracksdata/graph/_graph_view.py index be3eda2f..0b265fa4 100644 --- a/src/tracksdata/graph/_graph_view.py +++ b/src/tracksdata/graph/_graph_view.py @@ -1117,51 +1117,102 @@ def _split_node_attr_keys(self, attr_keys: Sequence[str]) -> tuple[list[str], li self._root._validate_attr_keys(root_keys, "node") return [k for k in attr_keys if k in local], root_keys - def _node_attrs_from_node_ids( + def _local_node_attrs_from_node_ids( self, + id_keys: list[str], *, node_ids: list[int] | None = None, attr_keys: Sequence[str] | str | None = None, unpack: bool = False, ) -> pl.DataFrame: - id_keys = [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET] + """ + Read node attributes entirely from this view's own local storage. - if attr_keys is None: - # "every key" means every key this view holds -- a fallback would - # silently pull in the columns the view was built to leave out. - node_dfs = super()._node_attrs_from_node_ids( - node_ids=self._map_to_local(node_ids), - attr_keys=None, - unpack=unpack, - ) - return self._map_df_to_external(node_dfs, id_keys) + The plain, no-fallback case: every requested key (or every key this view + holds, if `attr_keys` is `None`) is assumed to already live in this + view's `rx_graph`, so this is a direct call into the base + implementation plus the local -> external ID remap every `GraphView` + read needs. - if isinstance(attr_keys, str): - attr_keys = [attr_keys] - attr_keys = list(dict.fromkeys(attr_keys)) + Parameters + ---------- + id_keys : list[str] + ID columns to remap from local to external IDs via + `_map_df_to_external` (NODE_ID, EDGE_SOURCE, EDGE_TARGET). + node_ids : list[int] | None + External node IDs to read, mapped to local IDs before the read. + `None` means every node in the view. + attr_keys : Sequence[str] | str | None + Keys to read, assumed local. `None` means every key this view holds. + unpack : bool + Whether to unpack array-typed columns into separate columns. - local_keys, root_keys = self._split_node_attr_keys(attr_keys) + Returns + ------- + pl.DataFrame + One row per requested node, with external IDs. + """ + node_dfs = super()._node_attrs_from_node_ids( + node_ids=self._map_to_local(node_ids), + attr_keys=attr_keys, + unpack=unpack, + ) + return self._map_df_to_external(node_dfs, id_keys) - if not root_keys: - node_dfs = super()._node_attrs_from_node_ids( - node_ids=self._map_to_local(node_ids), - attr_keys=attr_keys, - unpack=unpack, - ) - return self._map_df_to_external(node_dfs, id_keys) + def _fallback_node_attrs_from_node_ids( + self, + id_keys: list[str], + *, + node_ids: list[int] | None, + attr_keys: list[str], + local_keys: list[str], + root_keys: list[str], + unpack: bool, + ) -> pl.DataFrame: + """ + Read a mix of locally-held and root-only attribute keys, joined on NODE_ID. + + Only called when `root_fallback=True` and the request actually includes + at least one key this view does not hold locally. Fetches `local_keys` + from this view and `root_keys` from `self._root`, for exactly the rows + this view covers, and joins the two on NODE_ID. + Parameters + ---------- + id_keys : list[str] + ID columns to remap from local to external IDs via + `_map_df_to_external` (NODE_ID, EDGE_SOURCE, EDGE_TARGET). + node_ids : list[int] | None + External node IDs to read. `None` means every node in the view. + attr_keys : list[str] + The full requested key list (local and root keys combined), in the + order the result should be returned in. + local_keys : list[str] + The subset of `attr_keys` this view holds locally. + root_keys : list[str] + The subset of `attr_keys` this view does not hold; fetched from + `self._root` instead. + unpack : bool + Whether to unpack array-typed columns into separate columns. + + Returns + ------- + pl.DataFrame + One row per requested node, columns in `attr_keys` order, with + external IDs. + """ node_id_key = DEFAULT_ATTR_KEYS.NODE_ID external_ids = self.node_ids() if node_ids is None else list(node_ids) # NODE_ID is needed as the join key even when the caller did not ask for # it, and `unpack` is deferred until after the join so that it sees the # same columns it would have on a view holding all of them. - local_df = super()._node_attrs_from_node_ids( - node_ids=self._map_to_local(external_ids), + local_df = self._local_node_attrs_from_node_ids( + id_keys, + node_ids=external_ids, attr_keys=[node_id_key, *local_keys], unpack=False, ) - local_df = self._map_df_to_external(local_df, id_keys) # Ask the root for exactly the rows this view covers -- never for all of # its rows, which include the ones this view deliberately dropped. @@ -1177,6 +1228,66 @@ def _node_attrs_from_node_ids( return node_dfs + def _node_attrs_from_node_ids( + self, + *, + node_ids: list[int] | None = None, + attr_keys: Sequence[str] | str | None = None, + unpack: bool = False, + ) -> pl.DataFrame: + """ + Get the attributes of the nodes as a polars DataFrame. + + Dispatches to one of two strategies, so that a request never pays for + root access it does not need: `_local_node_attrs_from_node_ids` when + every requested key lives in this view (always true when + `root_fallback=False`, or when the request simply does not name an + excluded key), `_fallback_node_attrs_from_node_ids` only when at least + one requested key is held solely by the root. + + Parameters + ---------- + node_ids : list[int] | None + The external IDs of the nodes to get the attributes for. + If None, all nodes in the view are used. + attr_keys : Sequence[str] | str | None + The attribute keys to get. If None, every key this view holds is + used -- never a key only reachable via `root_fallback`, since that + would silently widen the result past what the view was built with. + unpack : bool + Whether to unpack array attributes into multiple scalar attributes. + + Returns + ------- + pl.DataFrame + A polars DataFrame with the attributes of the nodes, in `attr_keys` + order when `attr_keys` is given, with external IDs. + """ + id_keys = [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET] + + if attr_keys is None: + # "every key" means every key this view holds -- a fallback would + # silently pull in the columns the view was built to leave out. + return self._local_node_attrs_from_node_ids(id_keys, node_ids=node_ids, attr_keys=None, unpack=unpack) + + if isinstance(attr_keys, str): + attr_keys = [attr_keys] + attr_keys = list(dict.fromkeys(attr_keys)) + + local_keys, root_keys = self._split_node_attr_keys(attr_keys) + + if not root_keys: + return self._local_node_attrs_from_node_ids(id_keys, node_ids=node_ids, attr_keys=attr_keys, unpack=unpack) + + return self._fallback_node_attrs_from_node_ids( + id_keys, + node_ids=node_ids, + attr_keys=attr_keys, + local_keys=local_keys, + root_keys=root_keys, + unpack=unpack, + ) + def node_attrs( self, *, From 09e9cf6badabc76c24cf03d70d3aa2788e6d76ec Mon Sep 17 00:00:00 2001 From: Caroline Malin-Mayor Date: Sat, 17 Oct 2026 09:46:27 -0400 Subject: [PATCH 5/6] Remove validation from attr splitting function --- src/tracksdata/graph/_graph_view.py | 60 +++++++++++++++++++---------- 1 file changed, 39 insertions(+), 21 deletions(-) diff --git a/src/tracksdata/graph/_graph_view.py b/src/tracksdata/graph/_graph_view.py index 0b265fa4..716ca7b7 100644 --- a/src/tracksdata/graph/_graph_view.py +++ b/src/tracksdata/graph/_graph_view.py @@ -1097,25 +1097,39 @@ def predecessors( raise RuntimeError("Out of sync graph view cannot be used to get predecessors") return super().predecessors(node_ids, attr_keys, return_attrs=return_attrs) - def _split_node_attr_keys(self, attr_keys: Sequence[str]) -> tuple[list[str], list[str]]: + def _split_node_attr_keys(self, attr_keys: Sequence[str] | None) -> tuple[list[str], list[str]]: """ - Partition requested node attribute keys into those this view holds and - those only the root holds. + Partition requested node attribute keys by whether this view holds them + locally. - Only ever returns root keys when the view was built with - ``root_fallback=True``; otherwise every key is reported as local and the - usual validation in the local read path rejects the ones that are not. - Keys that exist on neither the view nor the root still raise, here via - the root's own validation. + Purely a membership check against this view's own `node_attr_keys()` -- + it does not consult `root_fallback` or validate `request_from_root` + against the root. Callers decide what to do with each half. + + `None` means "every key this view holds" -- it never reaches into the + root, since a fallback would silently pull in the columns the view was + built to leave out. + + Parameters + ---------- + attr_keys : Sequence[str] | None + The requested keys to partition. `None` means "every key this view + holds" (see above). + + Returns + ------- + tuple[list[str], list[str]] + `(present_locally, request_from_root)`, each in the order keys + appeared in `attr_keys` (or in `node_attr_keys()` order, when + `attr_keys` is `None`). """ - if not self._root_fallback: - return list(attr_keys), [] + if attr_keys is None: + return self.node_attr_keys(return_ids=True), [] local = set(self.node_attr_keys(return_ids=True)) - root_keys = [k for k in attr_keys if k not in local] - if root_keys: - self._root._validate_attr_keys(root_keys, "node") - return [k for k in attr_keys if k in local], root_keys + present_locally = [k for k in attr_keys if k in local] + request_from_root = [k for k in attr_keys if k not in local] + return present_locally, request_from_root def _local_node_attrs_from_node_ids( self, @@ -1265,20 +1279,24 @@ def _node_attrs_from_node_ids( """ id_keys = [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET] - if attr_keys is None: - # "every key" means every key this view holds -- a fallback would - # silently pull in the columns the view was built to leave out. - return self._local_node_attrs_from_node_ids(id_keys, node_ids=node_ids, attr_keys=None, unpack=unpack) - if isinstance(attr_keys, str): attr_keys = [attr_keys] - attr_keys = list(dict.fromkeys(attr_keys)) + elif attr_keys is not None: + attr_keys = list(dict.fromkeys(attr_keys)) + + if not self._root_fallback: + # Nothing to route: every key is either local (fine) or missing, + # and the usual local-read validation rejects the latter. + return self._local_node_attrs_from_node_ids(id_keys, node_ids=node_ids, attr_keys=attr_keys, unpack=unpack) local_keys, root_keys = self._split_node_attr_keys(attr_keys) if not root_keys: - return self._local_node_attrs_from_node_ids(id_keys, node_ids=node_ids, attr_keys=attr_keys, unpack=unpack) + return self._local_node_attrs_from_node_ids(id_keys, node_ids=node_ids, attr_keys=local_keys, unpack=unpack) + # root_keys is only ever non-empty when attr_keys was an explicit list -- + # _split_node_attr_keys(None) always returns an empty second element. + self._root._validate_attr_keys(root_keys, "node") return self._fallback_node_attrs_from_node_ids( id_keys, node_ids=node_ids, From 02cb69a3f2945af99d71c7ccc6b372002f7c7844 Mon Sep 17 00:00:00 2001 From: Caroline Malin-Mayor Date: Sat, 17 Oct 2026 09:47:49 -0400 Subject: [PATCH 6/6] Use same splitting logic in validation error message --- src/tracksdata/graph/_graph_view.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/tracksdata/graph/_graph_view.py b/src/tracksdata/graph/_graph_view.py index 716ca7b7..35e643e0 100644 --- a/src/tracksdata/graph/_graph_view.py +++ b/src/tracksdata/graph/_graph_view.py @@ -414,8 +414,8 @@ def _validate_attr_keys( raise if isinstance(attr_keys, str): attr_keys = [attr_keys] - missing = set(attr_keys) - set(self.node_attr_keys(return_ids=True)) - on_root = sorted(missing & set(self._root.node_attr_keys(return_ids=True))) + _, missing = self._split_node_attr_keys(attr_keys) + on_root = sorted(set(missing) & set(self._root.node_attr_keys(return_ids=True))) if not on_root: raise raise KeyError(