diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 64f4c90e..ff0f93b8 100755 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/benchmarks/graph_mutations.py b/benchmarks/graph_mutations.py index b20579cf..073201fb 100644 --- a/benchmarks/graph_mutations.py +++ b/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]() diff --git a/benchmarks/tracklet_nodes.py b/benchmarks/tracklet_nodes.py index 27662052..6e10fc9f 100644 --- a/benchmarks/tracklet_nodes.py +++ b/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 diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index b2f66b46..2a3af63c 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/src/tracksdata/graph/_sql_graph.py @@ -1,9 +1,11 @@ import binascii import functools +import operator import re +import sqlite3 import uuid import weakref -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterator, Sequence from enum import Enum from pathlib import Path from typing import TYPE_CHECKING, Any, TypeVar @@ -50,6 +52,75 @@ T = TypeVar("T") +# SQLite's defaults are tuned for small databases where durability matters more +# than speed: a 2 MB page cache, a rollback journal, and an fsync on every +# commit. At tracking scale both hurt badly -- the page cache means a near-total +# miss rate on every B-tree descent, and the per-commit fsync dominates the +# latency of interactive single-node edits. +# +# ``synchronous=NORMAL`` under WAL is crash-safe against process death; only an +# OS crash or power loss can lose the most recent commits. That is the right +# trade for derived tracking data, but it *is* a durability change, so pass +# ``sqlite_pragmas={"synchronous": "FULL"}`` (or ``{}`` to disable all of these) +# to opt out. +DEFAULT_SQLITE_PRAGMAS: dict[str, Any] = { + "journal_mode": "WAL", # commit without writing/fsyncing a rollback journal + "synchronous": "NORMAL", # fsync at checkpoint rather than at every commit + "cache_size": -262_144, # negative => KiB, i.e. 256 MB of page cache + "temp_store": "MEMORY", # keep sorter/scratch spill out of the filesystem + "mmap_size": 1 << 30, # 1 GiB memory-mapped read window +} + +# Conservative floor: SQLITE_MAX_VARIABLE_NUMBER before SQLite 3.32 (2020). +_LEGACY_SQLITE_MAX_VARIABLES = 999 +# SQLITE_MAX_VARIABLE_NUMBER default since SQLite 3.32. +_MODERN_SQLITE_MAX_VARIABLES = 32_766 + + +@functools.cache +def _sqlite_max_variables() -> int: + """Return the bound-variable ceiling of the linked SQLite library. + + SQLite raised ``SQLITE_MAX_VARIABLE_NUMBER`` from 999 to 32766 in 3.32 + (2020). Assuming the legacy limit forces write batches and ``IN (...)`` + lists roughly 30x smaller than necessary, which in turn pushes + :class:`_SQLIDSet` onto its on-disk scratch-table path for id lists that + would comfortably fit inline. + """ + try: + with sqlite3.connect(":memory:") as conn: + limit = conn.getlimit(sqlite3.SQLITE_LIMIT_VARIABLE_NUMBER) + except (AttributeError, sqlite3.Error): + # ``getlimit`` needs Python 3.11+; fall back to the version heuristic. + limit = _MODERN_SQLITE_MAX_VARIABLES if sqlite3.sqlite_version_info >= (3, 32) else _LEGACY_SQLITE_MAX_VARIABLES + # Leave headroom for the handful of non-id parameters a statement also binds. + return max(_LEGACY_SQLITE_MAX_VARIABLES, limit - 100) + + +def _register_sqlite_pragmas(engine: sa.Engine, pragmas: dict[str, Any]) -> None: + """Apply *pragmas* to every new connection handed out by *engine*. + + Registered on ``connect`` rather than executed once, because each pooled + connection carries its own cache/mmap settings. Failures are logged and + swallowed: ``journal_mode=WAL`` is rejected on some network filesystems and + on ``:memory:`` databases, and neither should be fatal. + """ + if not pragmas: + return + + @sa.event.listens_for(engine, "connect") + def _set_pragmas(dbapi_connection: Any, _connection_record: Any) -> None: + cursor = dbapi_connection.cursor() + try: + for name, value in pragmas.items(): + try: + cursor.execute(f"PRAGMA {name}={value}") + except Exception as exc: # pragma: no cover - filesystem dependent + LOG.debug("Failed to set SQLite PRAGMA %s=%s: %s", name, value, exc) + finally: + cursor.close() + + def _is_builtin(obj: Any) -> bool: """Check if an object is a built-in type.""" return getattr(obj.__class__, "__module__", None) == "builtins" @@ -92,21 +163,69 @@ def _resolve_attr_filter_column( return getattr(table, flat_col) -def _to_sql_clause(f: Filter, table: type[DeclarativeBase]) -> Any: +def _time_band_clause( + table: type[DeclarativeBase], + f: AttrComparison, + multiplier: int, +) -> Any | None: + """Return a ``node_id`` range clause implied by a ``t == k`` comparison. + + Node ids assigned by :meth:`SQLGraph.bulk_add_nodes` are + ``t * multiplier + n``, so every node at time ``k`` necessarily has an id in + ``[k * multiplier, (k + 1) * multiplier)``. The returned clause is therefore + *logically redundant* -- it is implied by ``t == k`` and cannot change which + rows match under any boolean composition. It exists purely so the planner + has an indexed range to seek on: ``t`` is unindexed, so without it + ``WHERE t = k`` degrades to a full table scan. + + Returns ``None`` when the comparison is not an integer equality on ``t``, or + when the table has no ``node_id`` column (e.g. the edge table). + """ + if f.op is not operator.eq or f.attr.field_path or str(f.column) != DEFAULT_ATTR_KEYS.T: + return None + + other = f.other + # bool is an int subclass but is never a meaningful time point. + if isinstance(other, bool) or not isinstance(other, int | np.integer): + return None + + node_id_col = getattr(table, DEFAULT_ATTR_KEYS.NODE_ID, None) + if node_id_col is None: + return None + + time = int(other) + return sa.and_(node_id_col >= time * multiplier, node_id_col < (time + 1) * multiplier) + + +def _to_sql_clause( + f: Filter, + table: type[DeclarativeBase], + time_multiplier: int | None = None, +) -> Any: """Translate an AttrComparison or AttrFilter into a SQLAlchemy clause. Routes ``AttrComparison`` leaves through ``_resolve_attr_filter_column`` so struct-field comparisons resolve to the flat physical column. + + When *time_multiplier* is given, ``t == k`` leaves are additionally + constrained by the equivalent ``node_id`` range -- see + :func:`_time_band_clause`. Pass ``None`` to disable, which callers must do + whenever node ids are not known to encode their time point. """ if isinstance(f, AttrComparison): - return f.op(_resolve_attr_filter_column(table, f), f.other) + clause = f.op(_resolve_attr_filter_column(table, f), f.other) + if time_multiplier is not None: + band = _time_band_clause(table, f, time_multiplier) + if band is not None: + clause = sa.and_(clause, band) + return clause assert isinstance(f, AttrFilter) if f.op == "not": # AttrFilter.__init__ enforces exactly one operand for "not" - return sa.not_(_to_sql_clause(f.operands[0], table)) + return sa.not_(_to_sql_clause(f.operands[0], table, time_multiplier)) - clauses = [_to_sql_clause(o, table) for o in f.operands] + clauses = [_to_sql_clause(o, table, time_multiplier) for o in f.operands] if f.op == "and": return sa.and_(*clauses) if f.op == "or": @@ -194,6 +313,7 @@ def _filter_query( query: sa.Select, table: type[DeclarativeBase], attr_filters: Sequence[Filter], + time_multiplier: int | None = None, ) -> sa.Select: """ Filter a query by a list of attribute filters (AND-ed together at the top @@ -208,6 +328,10 @@ def _filter_query( The table to filter. attr_filters : Sequence[Filter] The attribute filters to apply. + time_multiplier : int | None + When given, ``t == k`` comparisons also emit the equivalent ``node_id`` + range so the planner can seek instead of scanning. See + :func:`_time_band_clause`. Returns ------- @@ -215,7 +339,7 @@ def _filter_query( The filtered query. """ LOG.info("Filter query:\n%s", attr_filters) - query = query.filter(*[_to_sql_clause(f, table) for f in attr_filters]) + query = query.filter(*[_to_sql_clause(f, table, time_multiplier) for f in attr_filters]) return query @@ -268,6 +392,10 @@ def __init__( self._edge_query = self._edge_query.filter(id_set.in_clause(self._graph.Edge.source_id)) node_filtered = True + # Only valid while node ids still encode their time point; the graph + # turns this off permanently if a custom index ever breaks the encoding. + time_multiplier = self._graph.node_id_time_multiplier if self._graph._time_encoded_ids else None + if self._node_attr_comps: node_filtered = True # filtering nodes by attributes @@ -275,6 +403,7 @@ def __init__( self._node_query, self._graph.Node, self._node_attr_comps, + time_multiplier, ) # if both node and edge attributes are filtered @@ -295,6 +424,7 @@ def __init__( self._edge_query, SourceNode, self._node_attr_comps, + time_multiplier, ) if self._include_sources or include_none: @@ -306,6 +436,7 @@ def __init__( self._edge_query, TargetNode, self._node_attr_comps, + time_multiplier, ) if self._edge_attr_comps: @@ -517,20 +648,28 @@ def subgraph( nodes_df = self._read_attr_dataframe(node_query, self._graph.Node) edges_df = self._read_attr_dataframe(edge_query, self._graph.Edge) - node_map_to_root = {} - node_map_from_root = {} rx_graph = rx.PyDiGraph() - for data in nodes_df.iter_rows(named=True): - root_node_id = data.pop(DEFAULT_ATTR_KEYS.NODE_ID) - node_id = rx_graph.add_node(data) - node_map_to_root[node_id] = root_node_id - node_map_from_root[root_node_id] = node_id - - for data in edges_df.iter_rows(named=True): - source_id = node_map_from_root[data.pop(DEFAULT_ATTR_KEYS.EDGE_SOURCE)] - target_id = node_map_from_root[data.pop(DEFAULT_ATTR_KEYS.EDGE_TARGET)] - rx_graph.add_edge(source_id, target_id, data) + # Build the view with two bulk rustworkx calls rather than one call per + # row. A frame-sized subgraph is tens of thousands of nodes and edges, + # and the per-row `add_node` / `add_edge` round trips dominated the cost + # of materializing it. + root_node_ids = nodes_df[DEFAULT_ATTR_KEYS.NODE_ID].to_list() + node_payloads = nodes_df.drop(DEFAULT_ATTR_KEYS.NODE_ID).to_dicts() + local_node_ids = list(rx_graph.add_nodes_from(node_payloads)) + + node_map_to_root = dict(zip(local_node_ids, root_node_ids, strict=True)) + node_map_from_root = dict(zip(root_node_ids, local_node_ids, strict=True)) + + edge_sources = edges_df[DEFAULT_ATTR_KEYS.EDGE_SOURCE].to_list() + edge_targets = edges_df[DEFAULT_ATTR_KEYS.EDGE_TARGET].to_list() + edge_payloads = edges_df.drop(DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET).to_dicts() + rx_graph.add_edges_from( + [ + (node_map_from_root[source_id], node_map_from_root[target_id], data) + for source_id, target_id, data in zip(edge_sources, edge_targets, edge_payloads, strict=True) + ] + ) graph = GraphView( rx_graph=rx_graph, @@ -579,9 +718,18 @@ class SQLGraph(BaseGraph): Database host. Not required for SQLite. port : int, optional Database port. Not required for SQLite. + engine_kwargs : dict[str, Any], optional + Extra keyword arguments forwarded to ``sqlalchemy.create_engine``. overwrite : bool, default False If True, drop and recreate all tables. Use with caution as this will delete all existing data. + sqlite_pragmas : dict[str, Any], optional + PRAGMAs applied to every SQLite connection. Defaults to + [DEFAULT_SQLITE_PRAGMAS][tracksdata.graph._sql_graph.DEFAULT_SQLITE_PRAGMAS], + which trades a small amount of durability for large gains in write + latency and cache hit rate — see that constant for the exact trade. + Pass ``{}`` to keep SQLite's own defaults, or a dict to override + individual PRAGMAs. Ignored for non-SQLite drivers. Attributes ---------- @@ -644,6 +792,7 @@ def __init__( port: int | None = None, engine_kwargs: dict[str, Any] | None = None, overwrite: bool = False, + sqlite_pragmas: dict[str, Any] | None = None, ): self._url = sa.engine.URL.create( drivername, @@ -654,7 +803,13 @@ def __init__( database=database, ) self._engine_kwargs = engine_kwargs if engine_kwargs is not None else {} - self._engine = sa.create_engine(self._url, **self._engine_kwargs) + self._sqlite_pragmas = sqlite_pragmas + self._engine = self._create_engine() + + # Whether every node id encodes its time point as + # ``t * node_id_time_multiplier + n``. Recomputed from the data in + # ``_update_max_id_per_time`` and maintained by ``bulk_add_nodes``. + self._time_encoded_ids = True # Create unique classes for this instance self._define_schema(overwrite=overwrite) @@ -669,6 +824,14 @@ def __init__( self._node_attr_schemas_cache: dict | None = None self._edge_attr_schemas_cache: dict | None = None + def _create_engine(self) -> sa.Engine: + """Build the engine and attach dialect-specific connection tuning.""" + engine = sa.create_engine(self._url, **self._engine_kwargs) + if engine.dialect.name == "sqlite": + pragmas = DEFAULT_SQLITE_PRAGMAS if self._sqlite_pragmas is None else self._sqlite_pragmas + _register_sqlite_pragmas(engine, pragmas) + return engine + def supports_custom_indices(self) -> bool: return True @@ -704,24 +867,44 @@ class Base(DeclarativeBase): class Node(Base): __tablename__ = "Node" - # Use node_id as sole primary key for simpler updates - node_id = sa.Column(sa.BigInteger, primary_key=True, unique=True) + # Use node_id as sole primary key for simpler updates. + # + # The SQLite variant is deliberately ``INTEGER`` rather than + # ``BIGINT``: SQLite only aliases a primary key onto the rowid when + # the declared type is literally ``INTEGER``. With ``BIGINT`` the + # table gets a hidden rowid *plus* a separate unique index, so every + # node is stored twice and each lookup is an index seek followed by + # a rowid seek. As the alias, ids assigned as + # ``t * node_id_time_multiplier + n`` also cluster the table + # physically by time, which is what makes a per-frame id range scan + # sequential. The value range is unchanged -- SQLite rowids are + # signed 64-bit. + node_id = sa.Column( + sa.BigInteger().with_variant(sa.Integer(), "sqlite"), + primary_key=True, + ) - # Add t as a regular column - # NOTE might want to use as index for fast time-based queries + # Add t as a regular column. Kept unindexed on purpose: `t == k` + # filters carry an equivalent `node_id` range (see + # `_time_band_clause`), so an index here would cost a write per node + # without speeding up the query it would serve. t = sa.Column(sa.Integer, nullable=False) node_tb_name = Node.__tablename__ + # NOTE: no `unique=True` on the primary keys below. It is implied by + # `primary_key=True`, and stating it again makes SQLAlchemy emit a + # redundant UNIQUE constraint -- an entire extra B-tree over the table, + # which on the edge table is the single largest avoidable index. class Edge(Base): __tablename__ = "Edge" - edge_id = sa.Column(sa.Integer, sa.Identity(always=True), primary_key=True, unique=True) + edge_id = sa.Column(sa.Integer, sa.Identity(always=True), primary_key=True) source_id = sa.Column(sa.BigInteger, sa.ForeignKey(f"{node_tb_name}.node_id"), index=True, nullable=False) target_id = sa.Column(sa.BigInteger, sa.ForeignKey(f"{node_tb_name}.node_id"), index=True, nullable=False) class Overlap(Base): __tablename__ = "Overlap" - overlap_id = sa.Column(sa.Integer, sa.Identity(always=True), primary_key=True, unique=True) + overlap_id = sa.Column(sa.Integer, sa.Identity(always=True), primary_key=True) source_id = sa.Column(sa.BigInteger, sa.ForeignKey(f"{node_tb_name}.node_id"), index=True, nullable=False) target_id = sa.Column(sa.BigInteger, sa.ForeignKey(f"{node_tb_name}.node_id"), index=True, nullable=False) @@ -941,10 +1124,33 @@ def _update_max_id_per_time(self) -> None: Scans the database to find the current maximum node ID for each time point and updates the internal cache to ensure newly created nodes have unique IDs. + + Also re-derives ``_time_encoded_ids`` from the data. Taking the per-time + minimum alongside the maximum costs nothing extra on top of the + group-by that is already needed, and deriving the flag rather than + trusting stored metadata means a database written by an older version -- + or via the raw SQL copy path in :meth:`_sqlite_table_dump` -- can never + silently enable the ``node_id`` range rewrite on ids that don't encode + their time point. """ with Session(self._engine) as session: - stmt = sa.select(self.Node.t, sa.func.max(self.Node.node_id)).group_by(self.Node.t) - self._max_id_per_time = {int(time): int(max_id) for time, max_id in session.execute(stmt).all()} + stmt = sa.select( + self.Node.t, + sa.func.min(self.Node.node_id), + sa.func.max(self.Node.node_id), + ).group_by(self.Node.t) + rows = [(int(time), int(min_id), int(max_id)) for time, min_id, max_id in session.execute(stmt).all()] + + self._max_id_per_time = {time: max_id for time, _, max_id in rows} + self._time_encoded_ids = all( + self._node_id_encodes_time(time, min_id) and self._node_id_encodes_time(time, max_id) + for time, min_id, max_id in rows + ) + + def _node_id_encodes_time(self, time: int, node_id: int) -> bool: + """Whether *node_id* falls in the id band that :meth:`bulk_add_nodes` assigns for *time*.""" + multiplier = self.node_id_time_multiplier + return time * multiplier <= node_id < (time + 1) * multiplier def filter( self, @@ -1039,6 +1245,12 @@ def bulk_add_nodes( self._max_id_per_time[time] = node_id else: node_id = indices[i] + # A caller-supplied id outside the time band breaks the + # assumption behind the `node_id` range rewrite in + # `_to_sql_clause`, so disable it for the rest of this graph's + # life rather than return wrong rows for `t == k` queries. + if self._time_encoded_ids and not self._node_id_encodes_time(time, node_id): + self._time_encoded_ids = False node_ids.append(node_id) insert_rows.append({**node, DEFAULT_ATTR_KEYS.NODE_ID: node_id}) @@ -1504,9 +1716,47 @@ def edge_ids(self) -> list[int]: return [i for (i,) in session.query(self.Edge.edge_id).all()] def time_points(self) -> list[int]: + if self._time_encoded_ids: + return self._time_points_by_id_bands() with Session(self._engine) as session: return [t for (t,) in session.query(self.Node.t).distinct().all()] + def _time_points_by_id_bands(self) -> list[int]: + """List the occupied time points by skipping between node id bands. + + ``SELECT DISTINCT t`` has to read every row -- and every row carries the + mask blob, so on a segmented graph that is the entire database. When ids + encode their time point, each occupied time point instead costs one + ``MIN(node_id)`` seek over the primary key: the walk jumps straight from + a frame's first node to the start of the next frame's band, so the cost + is proportional to the number of frames rather than the number of nodes. + + Termination is guaranteed by the same invariant that gates this path: a + node at time ``t`` has ``node_id < (t + 1) * multiplier``, so each step + seeks strictly past the current row and the walk stops when the seek + finds nothing. + """ + table = self.Node.__tablename__ + stmt = sa.text( + f""" + WITH RECURSIVE walk(tval) AS ( + SELECT t FROM "{table}" WHERE node_id = (SELECT MIN(node_id) FROM "{table}") + UNION ALL + SELECT ( + SELECT n.t FROM "{table}" AS n + WHERE n.node_id = ( + SELECT MIN(m.node_id) FROM "{table}" AS m + WHERE m.node_id >= (walk.tval + 1) * :multiplier + ) + ) + FROM walk WHERE walk.tval IS NOT NULL + ) + SELECT tval FROM walk WHERE tval IS NOT NULL + """ + ) + with Session(self._engine) as session: + return [int(t) for (t,) in session.execute(stmt, {"multiplier": self.node_id_time_multiplier}).all()] + def _reorder_by_indices( self, df: pl.DataFrame, @@ -2010,13 +2260,17 @@ def num_nodes(self) -> int: return int(session.query(self.Node).count()) def _sql_chunk_size(self) -> int: - if self._engine.dialect.name == "postgresql": - chunk_size = 30_000 - else: # for now everything else will use sqlite chunk size - # elif self._engine.dialect.name == "sqlite": - chunk_size = 900 + """Maximum number of bound values to put in a single statement. - return chunk_size + Write batches are chunked to this (divided by the column count), and it + also sets the cutoff above which :class:`_SQLIDSet` spills an id list + into an on-disk scratch table instead of inlining it. + """ + if self._engine.dialect.name == "postgresql": + return 30_000 + # Everything else is treated as SQLite. Query the linked library rather + # than assuming the pre-3.32 limit of 999. + return _sqlite_max_variables() def _update_table( self, @@ -2033,7 +2287,7 @@ def _update_table( if hasattr(ids, "tolist"): ids = ids.tolist() - # Handle array values with bulk_update_mappings + # Handle array values with a per-row bulk update attrs = attrs.copy() _data_numpy_to_native(attrs) schemas = self._attr_schemas_for_table(table_class) @@ -2072,7 +2326,46 @@ def _update_table( LOG.info("update %s table with %d rows", table_class.__table__, len(update_data)) LOG.info("update data sample: %s", update_data[:2]) - self._chunked_sa_write(Session.bulk_update_mappings, update_data, table_class) + self._bulk_update_by_id(table_class, id_key, update_data) + + # Bind parameter standing in for the row's primary key. Prefixed so it + # cannot collide with a user-defined attribute key, which becomes a real + # column name and therefore its own bind parameter below. + _PK_BIND_PARAM = "_tracksdata_pk" + + def _bulk_update_by_id( + self, + table_class: type[DeclarativeBase], + id_key: str, + rows: list[dict[str, Any]], + ) -> None: + """Apply per-row updates keyed by primary key, via a Core executemany. + + The ORM's ``bulk_update_mappings`` builds a unit-of-work update command + per row in Python, which is the dominant cost of any write-back over a + whole graph (``assign_tracklet_ids``, the solvers' solution columns). + Compiling the statement once and handing the DBAPI a parameter list lets + the driver run the loop instead, at roughly twice the throughput. + """ + if len(rows) == 0: + return + + table = table_class.__table__ + value_keys = [key for key in rows[0] if key != id_key] + if not value_keys: + return + + stmt = ( + sa.update(table) + .where(table.c[id_key] == sa.bindparam(self._PK_BIND_PARAM)) + # Binding through the column keeps each parameter's type -- notably + # PickleType -- so values are serialized exactly as on insert. + .values({key: sa.bindparam(key) for key in value_keys}) + ) + params = [{self._PK_BIND_PARAM: row[id_key], **{key: row[key] for key in value_keys}} for row in rows] + + with self._engine.begin() as conn: + conn.execute(stmt, params) def _chunked_sa_write( self, @@ -2183,33 +2476,41 @@ def update_node_attrs( if "t" in attrs: raise ValueError("Node attribute 't' cannot be updated.") + # Without a listener nothing needs the attribute payload, so skip the + # reads entirely. `node_ids=None` in particular must stay None so + # `_update_table` issues a single unqualified UPDATE rather than + # materializing every id in the graph just to filter on them. + if not is_signal_on(self.node_updated): + self._update_table(self.Node, node_ids, DEFAULT_ATTR_KEYS.NODE_ID, attrs) + return + updated_node_ids = self.node_ids() if node_ids is None else list(node_ids) if len(updated_node_ids) == 0: return attr_keys = self.node_attr_keys() - if is_signal_on(self.node_updated): - old_df = self.filter(node_ids=updated_node_ids).node_attrs( - attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, *attr_keys] - ) - old_attrs_by_id = old_df.rows_by_key( - key=DEFAULT_ATTR_KEYS.NODE_ID, named=True, unique=True, include_key=True - ) + old_df = self.filter(node_ids=updated_node_ids).node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, *attr_keys]) + old_attrs_by_id = old_df.rows_by_key(key=DEFAULT_ATTR_KEYS.NODE_ID, named=True, unique=True, include_key=True) self._update_table(self.Node, node_ids, DEFAULT_ATTR_KEYS.NODE_ID, attrs) - if is_signal_on(self.node_updated): - new_df = self.filter(node_ids=updated_node_ids).node_attrs( - attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, *attr_keys] - ) - new_attrs_by_id = new_df.rows_by_key( - key=DEFAULT_ATTR_KEYS.NODE_ID, named=True, unique=True, include_key=True - ) - 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()), - ) + # Derive the post-update payload from the pre-update one plus the values + # just written, rather than re-reading every attribute of every touched + # node. `_update_table` writes exactly `attrs` and nothing else, so the + # result is identical -- and it matches what `RustWorkXGraph` already + # emits, which reads back the written values rather than a round-trip. + n_nodes = len(updated_node_ids) + # Same broadcasting rule `_update_table` applies: scalars fan out across + # the batch, anything else is already one value per node. + written = {key: ([value] * n_nodes if np.isscalar(value) else list(value)) for key, value in attrs.items()} + + def _events() -> Iterator[tuple[int, dict[str, Any], dict[str, Any]]]: + for i, node_id in enumerate(updated_node_ids): + old_attrs = old_attrs_by_id[node_id] + new_attrs = {**old_attrs, **{key: values[i] for key, values in written.items()}} + yield node_id, old_attrs, new_attrs + + emit_node_updated_events(self.node_updated, _events(), set(attrs.keys())) def update_edge_attrs( self, @@ -2494,7 +2795,8 @@ def __getstate__(self) -> dict: def __setstate__(self, state: dict) -> None: self.__dict__.update(state) # recreate deleted objects - self._engine = sa.create_engine(self._url, **self._engine_kwargs) + self._sqlite_pragmas = state.get("_sqlite_pragmas") + self._engine = self._create_engine() self._define_schema(overwrite=False) def tracklet_graph( diff --git a/src/tracksdata/graph/_test/test_graph_backends.py b/src/tracksdata/graph/_test/test_graph_backends.py index 3c8095c6..a1a2ee8d 100644 --- a/src/tracksdata/graph/_test/test_graph_backends.py +++ b/src/tracksdata/graph/_test/test_graph_backends.py @@ -1840,6 +1840,42 @@ def test_sql_graph_mask_update_survives_reload(tmp_path: Path) -> None: np.testing.assert_array_equal(stored_mask.mask, mask_data) +def test_sql_graph_bulk_update_by_id_mixed_types(tmp_path: Path) -> None: + """Per-row updates go through a Core executemany; every column type must survive. + + The parameters are bound through their columns so that typed columns -- + notably ``PickleType`` for masks and arrays -- keep the same serialization + they get on insert. This pins that for a multi-row batch mixing pickled and + native columns, each row receiving its own value. + """ + graph = SQLGraph("sqlite", str(tmp_path / "bulk_update.db")) + 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("score", pl.Float64) + + node_ids = graph.bulk_add_nodes([{DEFAULT_ATTR_KEYS.T: t} for t in range(3)]) + + masks = [Mask(np.full((2, 2), i % 2 == 0, dtype=bool), bbox=np.array([i, i, i + 2, i + 2])) for i in range(3)] + bboxes = [np.array([i, i, i + 2, i + 2]) for i in range(3)] + graph.update_node_attrs( + node_ids=node_ids, + attrs={ + DEFAULT_ATTR_KEYS.MASK: masks, + DEFAULT_ATTR_KEYS.BBOX: bboxes, + "score": [1.5, 2.5, 3.5], + }, + ) + + df = graph.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.MASK, DEFAULT_ATTR_KEYS.BBOX, "score"]) + assert df["score"].to_list() == [1.5, 2.5, 3.5] + for i, (stored_mask, stored_bbox) in enumerate( + zip(df[DEFAULT_ATTR_KEYS.MASK].to_list(), df[DEFAULT_ATTR_KEYS.BBOX].to_list(), strict=True) + ): + assert isinstance(stored_mask, Mask) + np.testing.assert_array_equal(stored_mask.mask, masks[i].mask) + np.testing.assert_array_equal(np.asarray(stored_bbox), bboxes[i]) + + def test_sql_graph_struct_dtype_survives_reload(tmp_path: Path) -> None: db_path = tmp_path / "struct_graph.db" graph = SQLGraph("sqlite", str(db_path)) @@ -1874,6 +1910,232 @@ def test_sql_graph_max_id_restored_per_timepoint(tmp_path: Path) -> None: assert next_id == first_id + 1 +def _sqlite_query_plan(graph: SQLGraph, query: sa.Select) -> str: + """Return SQLite's EXPLAIN QUERY PLAN output for *query* as one string.""" + raw = graph._raw_query(query) + with graph._engine.connect() as conn: + rows = conn.execute(sa.text("EXPLAIN QUERY PLAN " + raw)).fetchall() + return "\n".join(str(row[-1]) for row in rows) + + +def _time_encoded_graph(db_path: Path, n_times: int = 4, per_time: int = 5) -> SQLGraph: + graph = SQLGraph("sqlite", str(db_path)) + graph.add_node_attr_key("score", pl.Float64) + for t in range(n_times): + graph.bulk_add_nodes([{DEFAULT_ATTR_KEYS.T: t, "score": float(i)} for i in range(per_time)]) + return graph + + +def test_sql_graph_time_filter_uses_node_id_range(tmp_path: Path) -> None: + """``t == k`` must seek on the node_id index instead of scanning the table. + + ``t`` is unindexed, so without the redundant node_id range conjunct every + per-frame query degrades to a full scan of the node table. + """ + graph = _time_encoded_graph(tmp_path / "time_band.db") + assert graph._time_encoded_ids + + filtered = graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == 2) + plan = _sqlite_query_plan(graph, filtered._node_query) + + # `node_id` is the rowid alias, so SQLite reports the seek as + # "INTEGER PRIMARY KEY (rowid>? AND rowid None: + """The node_id conjunct is logically implied, so it must not change results. + + Exercises the boolean compositions where a wrongly-scoped extra predicate + would show up: OR, NOT, and a plain conjunction with another attribute. + """ + graph = _time_encoded_graph(tmp_path / "compound.db") + + t_attr = NodeAttr(DEFAULT_ATTR_KEYS.T) + both_frames = graph.filter((t_attr == 1) | (t_attr == 3)).node_ids() + assert sorted(both_frames) == sorted(graph.filter(t_attr == 1).node_ids() + graph.filter(t_attr == 3).node_ids()) + + not_frame = graph.filter(~(t_attr == 1)).node_ids() + assert sorted(not_frame) == sorted(set(graph.node_ids()) - set(graph.filter(t_attr == 1).node_ids())) + + combined = graph.filter(t_attr == 2, NodeAttr("score") >= 3.0).node_ids() + assert sorted(combined) == sorted(graph.node_ids()[13:15]) + + # Non-integer right-hand sides must fall through unoptimized, not silently + # produce an empty band. + assert graph.filter(t_attr >= 2).node_ids() == graph.node_ids()[10:] + + +def test_sql_graph_custom_indices_disable_time_band(tmp_path: Path) -> None: + """Ids outside the time band must permanently disable the range rewrite.""" + db_path = tmp_path / "custom_ids.db" + graph = SQLGraph("sqlite", str(db_path)) + graph.bulk_add_nodes([{DEFAULT_ATTR_KEYS.T: 0}, {DEFAULT_ATTR_KEYS.T: 1}], indices=[7, 11]) + + assert not graph._time_encoded_ids + assert graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == 1).node_ids() == [11] + + plan = _sqlite_query_plan(graph, graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == 1)._node_query) + assert "node_id" not in plan, plan + + # The flag is derived from the stored ids, so a reopened database must not + # re-enable the rewrite on ids that do not encode their time point. + graph._engine.dispose() + reloaded = SQLGraph("sqlite", str(db_path)) + assert not reloaded._time_encoded_ids + assert reloaded.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == 1).node_ids() == [11] + + +def test_sql_graph_custom_indices_inside_band_keep_optimization(tmp_path: Path) -> None: + """Custom ids that still encode their time point should stay optimized.""" + graph = SQLGraph("sqlite", str(tmp_path / "in_band.db")) + multiplier = SQLGraph.node_id_time_multiplier + graph.bulk_add_nodes( + [{DEFAULT_ATTR_KEYS.T: 0}, {DEFAULT_ATTR_KEYS.T: 1}], + indices=[3, multiplier + 9], + ) + + assert graph._time_encoded_ids + assert graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == 1).node_ids() == [multiplier + 9] + + +def test_sql_graph_applies_sqlite_pragmas(tmp_path: Path) -> None: + """Default connections should be WAL + NORMAL with a large page cache.""" + graph = SQLGraph("sqlite", str(tmp_path / "pragmas.db")) + with graph._engine.connect() as conn: + assert conn.execute(sa.text("PRAGMA journal_mode")).scalar() == "wal" + assert conn.execute(sa.text("PRAGMA synchronous")).scalar() == 1 + assert conn.execute(sa.text("PRAGMA cache_size")).scalar() == -262_144 + + +def test_sql_graph_sqlite_pragmas_overridable(tmp_path: Path) -> None: + """``sqlite_pragmas`` must let callers opt out of the durability trade.""" + default_off = SQLGraph("sqlite", str(tmp_path / "no_pragmas.db"), sqlite_pragmas={}) + with default_off._engine.connect() as conn: + assert conn.execute(sa.text("PRAGMA journal_mode")).scalar() == "delete" + assert conn.execute(sa.text("PRAGMA synchronous")).scalar() == 2 + + strict = SQLGraph("sqlite", str(tmp_path / "strict.db"), sqlite_pragmas={"synchronous": "FULL"}) + with strict._engine.connect() as conn: + assert conn.execute(sa.text("PRAGMA synchronous")).scalar() == 2 + + +def test_sql_graph_ddl_has_no_redundant_indexes(tmp_path: Path) -> None: + """Primary keys must not carry a second UNIQUE index. + + ``primary_key=True`` already enforces uniqueness; adding ``unique=True`` + makes SQLAlchemy emit an extra constraint, and SQLite backs that with a + whole additional B-tree over the table. + """ + graph = _time_encoded_graph(tmp_path / "ddl.db") + with graph._engine.connect() as conn: + autoindexes = { + row[0] + for row in conn.execute( + sa.text("SELECT name FROM sqlite_master WHERE type='index' AND name LIKE 'sqlite_autoindex_%'") + ).fetchall() + } + # Metadata's text primary key legitimately needs one; Node and Edge must not. + assert not any(name.startswith(("sqlite_autoindex_Node", "sqlite_autoindex_Edge")) for name in autoindexes), ( + autoindexes + ) + + +def test_sql_graph_node_id_is_sqlite_rowid_alias(tmp_path: Path) -> None: + """``node_id`` must be declared INTEGER so SQLite aliases it onto the rowid. + + With ``BIGINT`` the table keeps a hidden rowid plus a separate unique index, + which doubles node storage and turns every lookup into two seeks. + """ + graph = _time_encoded_graph(tmp_path / "rowid.db") + with graph._engine.connect() as conn: + ddl = conn.execute(sa.text("SELECT sql FROM sqlite_master WHERE name='Node'")).scalar() + assert "node_id INTEGER" in ddl, ddl + + plan = _sqlite_query_plan(graph, graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == 2)._node_query) + assert "INTEGER PRIMARY KEY" in plan, plan + + +def test_sql_graph_time_points_skips_between_id_bands(tmp_path: Path) -> None: + """``time_points`` must agree with a DISTINCT scan, including after removals.""" + graph = _time_encoded_graph(tmp_path / "time_points.db", n_times=6, per_time=3) + assert graph._time_encoded_ids + assert graph.time_points() == [0, 1, 2, 3, 4, 5] + + # Emptying a frame must drop it from the result. + frame_two = graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == 2).node_ids() + graph.bulk_remove_nodes(frame_two) + assert graph.time_points() == [0, 1, 3, 4, 5] + + # Removing the first frame must not stop the walk before it starts. + graph.bulk_remove_nodes(graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == 0).node_ids()) + assert graph.time_points() == [1, 3, 4, 5] + + # And the fallback path must agree. + graph._time_encoded_ids = False + assert sorted(graph.time_points()) == [1, 3, 4, 5] + + +def test_sql_graph_time_points_handles_negative_and_empty(tmp_path: Path) -> None: + """The band walk must handle negative time points and an empty graph.""" + graph = SQLGraph("sqlite", str(tmp_path / "negative_t.db")) + assert graph.time_points() == [] + + graph.bulk_add_nodes([{DEFAULT_ATTR_KEYS.T: t} for t in (-3, -3, -1, 0, 4)]) + assert graph._time_encoded_ids + assert graph.time_points() == [-3, -1, 0, 4] + + +def test_sql_graph_chunk_size_uses_real_variable_limit() -> None: + """The bind-variable budget must reflect the linked SQLite, not the 999 era.""" + graph = SQLGraph("sqlite", ":memory:") + assert graph._sql_chunk_size() > 999 + + +def test_sql_graph_update_node_attrs_signal_payload(tmp_path: Path) -> None: + """The derived post-update payload must match a full re-read. + + ``update_node_attrs`` builds ``new_attrs`` from the pre-update row plus the + written values instead of re-querying, so this pins that the emitted payload + still carries every attribute with the updated values applied. + """ + graph = _time_encoded_graph(tmp_path / "signals.db", n_times=2, per_time=3) + events: list[tuple[list[int], list[dict], list[dict], set[str]]] = [] + graph.node_updated.connect(lambda *args: events.append(args)) + + targets = graph.node_ids()[:2] + graph.update_node_attrs(node_ids=targets, attrs={"score": [10.0, 20.0]}) + + (node_ids, old_attrs, new_attrs, changed_keys) = events[-1] + assert node_ids == targets + assert changed_keys == {"score"} + assert [a["score"] for a in old_attrs] == [0.0, 1.0] + assert [a["score"] for a in new_attrs] == [10.0, 20.0] + # Untouched attributes must still be present and unchanged. + assert [a[DEFAULT_ATTR_KEYS.T] for a in new_attrs] == [0, 0] + assert set(new_attrs[0]) == set(old_attrs[0]) + # And the payload must agree with what the database now holds. + stored = graph.filter(node_ids=targets).node_attrs(attr_keys=["score"]) + assert stored["score"].to_list() == [10.0, 20.0] + + +def test_sql_graph_update_all_nodes_without_listener_skips_id_materialization(tmp_path: Path) -> None: + """``node_ids=None`` and no listener must not enumerate every node id.""" + graph = _time_encoded_graph(tmp_path / "update_all.db", n_times=2, per_time=3) + + def _fail() -> list[int]: + raise AssertionError("node_ids() should not be called without a listener") + + graph.node_ids = _fail # type: ignore[method-assign] + graph.update_node_attrs(attrs={"score": 5.0}) + del graph.node_ids + + assert graph.node_attrs(attr_keys=["score"])["score"].to_list() == [5.0] * 6 + + def test_sql_graph_schema_defaults_survive_reload(tmp_path: Path) -> None: """Reloading a SQLGraph should preserve dtype and default schema metadata.""" db_path = tmp_path / "schema_defaults.db" diff --git a/src/tracksdata/utils/_dataframe.py b/src/tracksdata/utils/_dataframe.py index a6de0f17..bdff4dce 100644 --- a/src/tracksdata/utils/_dataframe.py +++ b/src/tracksdata/utils/_dataframe.py @@ -1,6 +1,5 @@ import cloudpickle import polars as pl -import polars.selectors as cs def unpack_array_attrs(df: pl.DataFrame) -> pl.DataFrame: @@ -33,6 +32,12 @@ def unpickle_bytes_columns(df: pl.DataFrame) -> pl.DataFrame: """ Unpickle bytes columns from the database. + The result is left as :class:`polars.Object`. Every caller pairs this with + ``SQLGraph._cast_columns``, which casts each pickled column to its declared + schema dtype, so narrowing the dtype here would only build an intermediate + that is immediately rebuilt -- and for a column of opaque payloads (masks) + the attempt materializes every value just to fail and be discarded. + Parameters ---------- df : pl.DataFrame @@ -43,11 +48,18 @@ def unpickle_bytes_columns(df: pl.DataFrame) -> pl.DataFrame: pl.DataFrame The DataFrame with the bytes columns unpickled. """ - df = df.map_columns(cs.binary(), lambda x: x.map_elements(cloudpickle.loads, return_dtype=pl.Object)) - for col, dtype in zip(df.columns, df.dtypes, strict=True): - if isinstance(dtype, pl.Object): - try: - df = df.with_columns(pl.Series(df[col].to_list()).alias(col)) - except Exception: - pass - return df + binary_cols = [name for name, dtype in df.schema.items() if dtype == pl.Binary] + if not binary_cols: + return df + + # A plain comprehension rather than `map_elements`: these are opaque Python + # objects either way, so routing them through the expression engine adds + # per-element dispatch without buying any vectorization. + return df.with_columns( + pl.Series( + name, + [None if value is None else cloudpickle.loads(value) for value in df[name]], + dtype=pl.Object, + ) + for name in binary_cols + )