WIP: Performance improvement for SQLGraph

.pre-commit-config.yaml

@@ -20,20 +20,10 @@ repos:
     hooks:
       - id: absolufy-imports
 
-  # sorting imports
-  - repo: https://github.com/pycqa/isort
-    rev: 5.12.0
-    hooks:
-      - id: isort
-        args: ["--profile", "black", "--filter-files"]
-
-  # formatting
-  - repo: https://github.com/psf/black
-    rev: 25.1.0
-    hooks:
-      - id: black
-        args: ["--line-length", "120"]
-
+  # linting, import sorting (ruff "I") and formatting
+  # NOTE: ruff-format is the single formatter here; do not add black/isort back,
+  # they disagree with it (magic trailing comma, string quotes, import grouping)
+  # and CI only runs `ruff check` + `ruff format --check`.
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
     rev: v0.16.0

benchmarks/graph_mutations.py

@@ -109,6 +109,53 @@ def time_filter_node_ids(self, backend_name: str, n_nodes: int) -> None:
         self.graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) >= 1).node_ids()
 
 
+class InteractiveLatencyBenchmark:
+    """Per-frame query latency on a graph with many time points.
+
+    The mutation benchmarks above use ~50 time points, where a full scan of the
+    node table and an indexed seek into a single frame cost about the same. The
+    interactive-editing budget is set by the opposite shape -- thousands of
+    frames, where `filter(t == k)` touches a small slice of a large table -- so
+    this parametrizes on frame count instead of raw node count.
+
+    Guards the `node_id`-range rewrite of `t == k` filters: without it these
+    degrade to a full table scan and grow linearly with the graph.
+    """
+
+    param_names = ("backend", "n_frames")
+    params = (tuple(BACKENDS), (500,) if IS_CI else (500, 2_000))
+
+    timeout = 600
+
+    # Nodes per frame; total nodes are `n_frames * NODES_PER_FRAME`.
+    NODES_PER_FRAME = 200
+
+    def setup(self, backend_name: str, n_frames: int) -> None:
+        self.graph = BACKENDS[backend_name]()
+        self.graph.add_node_attr_key("score", dtype=pl.Float64)
+        for t in range(n_frames):
+            self.graph.bulk_add_nodes([{DEFAULT_ATTR_KEYS.T: t, "score": 0.0} for _ in range(self.NODES_PER_FRAME)])
+        self.mid_frame = n_frames // 2
+        self.frame_ids = self.graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == self.mid_frame).node_ids()
+
+    def time_frame_node_ids(self, backend_name: str, n_frames: int) -> None:
+        self.graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == self.mid_frame).node_ids()
+
+    def time_frame_node_attrs(self, backend_name: str, n_frames: int) -> None:
+        self.graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == self.mid_frame).node_attrs(attr_keys=["score"])
+
+    def time_frame_subgraph(self, backend_name: str, n_frames: int) -> None:
+        self.graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == self.mid_frame).subgraph()
+
+    def time_filter_by_frame_node_ids(self, backend_name: str, n_frames: int) -> None:
+        # Drives the `_SQLIDSet` inline-vs-scratch-table cutoff: a frame's worth
+        # of ids must stay inline rather than spilling to an on-disk table.
+        self.graph.filter(node_ids=self.frame_ids).node_attrs(attr_keys=["score"])
+
+    def time_update_single_node(self, backend_name: str, n_frames: int) -> None:
+        self.graph.update_node_attrs(node_ids=self.frame_ids[:1], attrs={"score": 1.0})
+
+
 def _build_bbox_graph(backend_name: str, n_nodes: int) -> td.graph.BaseGraph:
     """Graph whose nodes carry a bbox, so a real BBoxSpatialFilter can index them."""
     graph = BACKENDS[backend_name]()

benchmarks/tracklet_nodes.py

@@ -1,6 +1,7 @@
 from itertools import pairwise
 
 import tracksdata as td
+from tracksdata.attrs import NodeAttr
 
 if __name__ == "__main__":
     from common import BACKENDS, IS_CI  # For local testing
@@ -45,6 +46,60 @@ def time_tracklet_nodes(self, backend_name: str, n_nodes: int, n_lineages: int)
         return self.graph.tracklet_nodes([self.node_ids[len(self.node_ids) // 2]])
 
 
+class TrackletSubgraphBenchmark:
+    """`assign_tracklet_ids` and materializing a single tracklet as a subgraph.
+
+    Each lineage is an unbranched chain spanning every frame, so it maps to
+    exactly one tracklet. Selecting one track in a viewer filters on
+    `tracklet_id` -- unlike `t`, that column is neither indexed nor encoded in
+    the node ids, so the filter is a full scan of the node table and grows with
+    the whole graph rather than with the track.
+
+    `assign_tracklet_ids` is the counterpart worst case: the SQL backend
+    implements it by materializing the *entire* graph as a `GraphView`, so it
+    is bounded by subgraph construction rather than by the tracklet logic.
+    """
+
+    param_names = ("backend", "n_frames", "n_lineages")
+    params = (tuple(BACKENDS), (200,) if IS_CI else (200, 1_000), (100,))
+
+    timeout = 600
+
+    def setup(self, backend_name: str, n_frames: int, n_lineages: int) -> None:
+        self.graph = BACKENDS[backend_name]()
+        prev_ids: list[int] = []
+        for t in range(n_frames):
+            ids = self.graph.bulk_add_nodes([{td.DEFAULT_ATTR_KEYS.T: t} for _ in range(n_lineages)])
+            if prev_ids:
+                self.graph.bulk_add_edges(
+                    [
+                        {td.DEFAULT_ATTR_KEYS.EDGE_SOURCE: s, td.DEFAULT_ATTR_KEYS.EDGE_TARGET: d}
+                        for s, d in zip(prev_ids, ids, strict=True)
+                    ]
+                )
+            prev_ids = ids
+
+        self.graph.assign_tracklet_ids()
+        tracklet_key = td.DEFAULT_ATTR_KEYS.TRACKLET_ID
+        tracklet_ids = self.graph.node_attrs(attr_keys=[tracklet_key])[tracklet_key].unique().to_list()
+        self.target_tracklet = tracklet_ids[len(tracklet_ids) // 2]
+
+    def time_assign_tracklet_ids(self, backend_name: str, n_frames: int, n_lineages: int) -> None:
+        self.graph.assign_tracklet_ids()
+
+    def time_filter_tracklet_node_ids(self, backend_name: str, n_frames: int, n_lineages: int) -> None:
+        self.graph.filter(NodeAttr(td.DEFAULT_ATTR_KEYS.TRACKLET_ID) == self.target_tracklet).node_ids()
+
+    def time_subgraph_one_tracklet(self, backend_name: str, n_frames: int, n_lineages: int) -> None:
+        self.graph.filter(NodeAttr(td.DEFAULT_ATTR_KEYS.TRACKLET_ID) == self.target_tracklet).subgraph()
+
+    def time_subgraph_one_tracklet_by_ids(self, backend_name: str, n_frames: int, n_lineages: int) -> None:
+        # Two-step form: resolve the ids first, then materialize by id. Isolates
+        # subgraph construction from the unindexed tracklet_id scan.
+        node_ids = self.graph.filter(NodeAttr(td.DEFAULT_ATTR_KEYS.TRACKLET_ID) == self.target_tracklet).node_ids()
+        self.graph.filter(node_ids=node_ids).subgraph()
+
+
 if __name__ == "__main__":
     import cProfile