Skip to content

accept numpy integers as node/edge ids #339

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
TeunHuijben wants to merge 1 commit into
base: main
Choose a base branch
from
Open

accept numpy integers as node/edge ids #339

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions src/tracksdata/graph/_base_graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
)
from tracksdata.utils._logging import LOG
from tracksdata.utils._multiprocessing import multiprocessing_apply
from tracksdata.utils._numpy_native import is_int_like, to_native

if TYPE_CHECKING:
import motile
Expand Down Expand Up @@ -2276,9 +2277,9 @@ def __getitem__(self, node_id: int) -> "NodeInterface":
NodeInterface
Interface for accessing the node's attributes.
"""
if not isinstance(node_id, int):
if not is_int_like(node_id):
raise ValueError(f"node_id must be an integer, found '{node_id}' of type {type(node_id)}")
return NodeInterface(self._graph, node_id)
return NodeInterface(self._graph, to_native(node_id))


class EdgesAccessor:
Expand Down Expand Up @@ -2308,9 +2309,9 @@ def __getitem__(self, edge_id: int) -> "EdgeInterface":
EdgeInterface
Interface for accessing the edge's attributes.
"""
if not isinstance(edge_id, int):
if not is_int_like(edge_id):
raise ValueError(f"edge_id must be an integer, found '{edge_id}' of type {type(edge_id)}")
return EdgeInterface(self._graph, edge_id)
return EdgeInterface(self._graph, to_native(edge_id))


class NodeInterface:
Expand Down
23 changes: 9 additions & 14 deletions src/tracksdata/graph/_graph_view.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from tracksdata.graph._rustworkx_graph import IndexedRXGraph, RustWorkXGraph, RXFilter
from tracksdata.graph.filters._indexed_filter import IndexRXFilter
from tracksdata.utils._dtypes import AttrSchema
from tracksdata.utils._numpy_native import is_int_like, to_native, to_native_list
from tracksdata.utils._signal import (
emit_node_added_events,
emit_node_removed_events,
Expand Down Expand Up @@ -474,10 +475,7 @@ def bulk_remove_nodes(self, node_ids: Sequence[int]) -> None:
ValueError
If any node_id does not exist in the graph.
"""
if hasattr(node_ids, "tolist"):
node_ids = node_ids.tolist()
else:
node_ids = list(node_ids)
node_ids = to_native_list(node_ids)
if len(node_ids) == 0:
return

Expand Down Expand Up @@ -701,10 +699,7 @@ def bulk_remove_edges(self, edge_ids: Sequence[int]) -> None:
ValueError
If any edge_id does not exist in the root graph.
"""
if hasattr(edge_ids, "tolist"):
edge_ids = edge_ids.tolist()
else:
edge_ids = list(edge_ids)
edge_ids = to_native_list(edge_ids)
if len(edge_ids) == 0:
return

Expand Down Expand Up @@ -865,8 +860,8 @@ def _get_neighbors(
single_node = False
if node_ids is None:
node_ids = self.node_ids()
elif isinstance(node_ids, int):
node_ids = [node_ids]
elif is_int_like(node_ids):
node_ids = [to_native(node_ids)]
single_node = True

local_node_ids = self._map_to_local(node_ids)
Expand Down Expand Up @@ -1074,8 +1069,8 @@ def in_degree(self, node_ids: list[int] | int | None = None) -> list[int] | int:
if node_ids is None:
node_ids = self.node_ids()
rx_graph = self.rx_graph
if isinstance(node_ids, int):
return rx_graph.in_degree(self._map_to_local(node_ids))
if is_int_like(node_ids):
return rx_graph.in_degree(self._map_to_local(to_native(node_ids)))
return [rx_graph.in_degree(self._map_to_local(node_id)) for node_id in node_ids]

def out_degree(self, node_ids: list[int] | int | None = None) -> list[int] | int:
Expand All @@ -1085,8 +1080,8 @@ def out_degree(self, node_ids: list[int] | int | None = None) -> list[int] | int
if node_ids is None:
node_ids = self.node_ids()
rx_graph = self.rx_graph
if isinstance(node_ids, int):
return rx_graph.out_degree(self._map_to_local(node_ids))
if is_int_like(node_ids):
return rx_graph.out_degree(self._map_to_local(to_native(node_ids)))
return [rx_graph.out_degree(self._map_to_local(node_id)) for node_id in node_ids]

def dividing_nodes(self) -> list[int]:
Expand Down
38 changes: 14 additions & 24 deletions src/tracksdata/graph/_rustworkx_graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from tracksdata.utils._dataframe import unpack_array_attrs
from tracksdata.utils._dtypes import AttrSchema, process_attr_key_args
from tracksdata.utils._logging import LOG
from tracksdata.utils._numpy_native import is_int_like, to_native, to_native_list
from tracksdata.utils._signal import (
emit_node_added_events,
emit_node_removed_events,
Expand Down Expand Up @@ -179,8 +180,8 @@ def __init__(
self._graph = graph
self._attr_comps = attr_comps

if node_ids is not None and hasattr(node_ids, "tolist"):
node_ids = node_ids.tolist()
if node_ids is not None:
node_ids = to_native_list(node_ids)

self._node_ids = node_ids
self._include_targets = include_targets
Expand Down Expand Up @@ -691,10 +692,7 @@ def bulk_remove_nodes(self, node_ids: Sequence[int]) -> None:
ValueError
If any node_id does not exist in the graph.
"""
if hasattr(node_ids, "tolist"):
node_ids = node_ids.tolist()
else:
node_ids = list(node_ids)
node_ids = to_native_list(node_ids)
if len(node_ids) == 0:
return

Expand Down Expand Up @@ -750,10 +748,7 @@ def bulk_remove_edges(self, edge_ids: Sequence[int]) -> None:
ValueError
If any edge_id does not exist in the graph.
"""
if hasattr(edge_ids, "tolist"):
edge_ids = edge_ids.tolist()
else:
edge_ids = list(edge_ids)
edge_ids = to_native_list(edge_ids)
if len(edge_ids) == 0:
return
self._bulk_remove_edges_local(edge_ids)
Expand Down Expand Up @@ -841,8 +836,8 @@ def _get_neighbors(
rx_graph = self.rx_graph
if node_ids is None:
node_ids = list(rx_graph.node_indices())
elif isinstance(node_ids, int):
node_ids = [node_ids]
elif is_int_like(node_ids):
node_ids = [to_native(node_ids)]
single_node = True

if not return_attrs and attr_keys is not None:
Expand Down Expand Up @@ -1434,7 +1429,7 @@ def assign_tracklet_ids(
"Often used from `graph.subgraph(edge_attr_filter={'solution': True})`"
) from e

# Converting to list of int for SQLGraph compatibility (See below)
# A list, not the numpy array, so the id remapping below can index into it
tracklet_ids = tracklet_ids.tolist()

# For the IndexedRXGraph, we need to map the track_node_ids to the external node ids
Expand Down Expand Up @@ -1463,9 +1458,7 @@ def assign_tracklet_ids(
tracklet_id_map = dict(
zip(tracklet_id_map[output_key + "_new"], tracklet_id_map[output_key], strict=True)
)
# Ensure that the result is a list of integers (using numpy integer causes issues with SQLGraph)
# Later on, we will make it safe to use numpy integers everywhere for updating attributes.
tracklet_ids = [int(tracklet_id_map.get(tid, tid)) for tid in tracklet_ids] # type: ignore
tracklet_ids = [tracklet_id_map.get(tid, tid) for tid in tracklet_ids] # type: ignore
# Update the value with the reused IDs
id_update_df = id_update_df.with_columns(pl.Series(output_key + "_new", tracklet_ids))

Expand All @@ -1490,8 +1483,8 @@ def in_degree(self, node_ids: list[int] | int | None = None) -> list[int] | int:
if node_ids is None:
node_ids = self.node_ids()
rx_graph = self.rx_graph
if isinstance(node_ids, int):
return rx_graph.in_degree(node_ids)
if is_int_like(node_ids):
return rx_graph.in_degree(to_native(node_ids))
return [rx_graph.in_degree(node_id) for node_id in node_ids]

def out_degree(self, node_ids: list[int] | int | None = None) -> list[int] | int:
Expand All @@ -1501,8 +1494,8 @@ def out_degree(self, node_ids: list[int] | int | None = None) -> list[int] | int
if node_ids is None:
node_ids = self.node_ids()
rx_graph = self.rx_graph
if isinstance(node_ids, int):
return rx_graph.out_degree(node_ids)
if is_int_like(node_ids):
return rx_graph.out_degree(to_native(node_ids))
return [rx_graph.out_degree(node_id) for node_id in node_ids]

def dividing_nodes(self) -> list[int]:
Expand Down Expand Up @@ -2032,10 +2025,7 @@ def bulk_remove_nodes(self, node_ids: Sequence[int]) -> None:
ValueError
If any node_id does not exist in the graph.
"""
if hasattr(node_ids, "tolist"):
node_ids = node_ids.tolist()
else:
node_ids = list(node_ids)
node_ids = to_native_list(node_ids)
if len(node_ids) == 0:
return

Expand Down
63 changes: 27 additions & 36 deletions src/tracksdata/graph/_sql_graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
sqlalchemy_type_to_polars_dtype,
)
from tracksdata.utils._logging import LOG
from tracksdata.utils._numpy_native import is_int_like, to_native, to_native_list
from tracksdata.utils._signal import (
emit_node_added_events,
emit_node_removed_events,
Expand All @@ -65,22 +66,15 @@ def _data_numpy_to_native(data: dict[str, Any]) -> None:
"""
Convert numpy scalars to native Python scalars in place.

Database drivers do not know about numpy scalar types. ``sqlite3``, for example,
falls back to the buffer protocol and stores ``np.int64(7)`` as its raw
little-endian byte buffer (a BLOB), silently corrupting a column declared as
``BIGINT``. Numpy floats and strings happen to survive because they subclass
their Python counterparts, which makes the corruption look selective.
See :func:`tracksdata.utils._numpy_native.to_native` for why drivers need this.

Parameters
----------
data : dict[str, Any]
The data to convert. Modified in place.
"""
for k, v in data.items():
# `np.generic` is the base class of every numpy scalar, and excludes
# (0-dim) arrays, which must be passed through untouched.
if isinstance(v, np.generic):
data[k] = v.item()
data[k] = to_native(v)


def _resolve_attr_filter_column(
Expand Down Expand Up @@ -113,7 +107,11 @@ def _to_sql_clause(f: Filter, table: type[DeclarativeBase]) -> Any:
struct-field comparisons resolve to the flat physical column.
"""
if isinstance(f, AttrComparison):
return f.op(_resolve_attr_filter_column(table, f), f.other)
# The compared value is bound as a SQL parameter, so numpy scalars (including
# the sequence of them that `is_in` takes) must be converted first.
other = f.other
other = to_native_list(other) if isinstance(other, list | tuple | np.ndarray) else to_native(other)
return f.op(_resolve_attr_filter_column(table, f), other)

assert isinstance(f, AttrFilter)
if f.op == "not":
Expand Down Expand Up @@ -167,9 +165,7 @@ def __init__(
*,
occurrences: int = 1,
) -> None:
if hasattr(ids, "tolist"):
ids = ids.tolist()
self._ids: list[int] = list(ids)
self._ids: list[int] = to_native_list(ids)
# Hold the engine, not the graph, so this set does not participate in
# the graph -> SQLFilter -> _SQLIDSet -> graph reference cycle.
# Otherwise the scratch table would only be dropped after Python's
Expand Down Expand Up @@ -1085,10 +1081,7 @@ def bulk_remove_nodes(self, node_ids: Sequence[int]) -> None:
ValueError
If any node_id does not exist in the graph.
"""
if hasattr(node_ids, "tolist"):
node_ids = node_ids.tolist()
else:
node_ids = list(node_ids)
node_ids = to_native_list(node_ids)

if len(node_ids) == 0:
return
Expand Down Expand Up @@ -1240,10 +1233,10 @@ def bulk_add_overlaps(
[add_overlap][tracksdata.graph.SQLGraph.add_overlap]:
Add a single overlap to the graph.
"""
if hasattr(overlaps, "tolist"):
overlaps = overlaps.tolist()

overlaps = [{"source_id": int(source_id), "target_id": int(target_id)} for source_id, target_id in overlaps]
# `overlaps` is a nested sequence, so each id is converted individually
overlaps = [
{"source_id": to_native(source_id), "target_id": to_native(target_id)} for source_id, target_id in overlaps
]
self._chunked_sa_write(Session.bulk_insert_mappings, overlaps, self.Overlap)

def overlaps(
Expand All @@ -1260,8 +1253,8 @@ def overlaps(
filtered in Polars afterwards to avoid a quadratic blow-up of bound
parameters.
"""
if hasattr(node_ids, "tolist"):
node_ids = node_ids.tolist()
if node_ids is not None:
node_ids = to_native_list(node_ids)

with Session(self._engine) as session:
base_query = session.query(self.Overlap.source_id, self.Overlap.target_id)
Expand Down Expand Up @@ -1330,14 +1323,15 @@ def _get_neighbors(
"""
single_node = False
filter_node_ids: list[int] | None
if isinstance(node_ids, int):
node_ids = [node_ids]
if is_int_like(node_ids):
node_ids = [to_native(node_ids)]
filter_node_ids = node_ids
single_node = True
elif node_ids is None:
node_ids = self.node_ids()
filter_node_ids = None
else:
node_ids = to_native_list(node_ids)
filter_node_ids = node_ids

if isinstance(attr_keys, str):
Expand Down Expand Up @@ -2036,8 +2030,7 @@ def _update_table(
LOG.info("No ids to update, skipping")
return

if hasattr(ids, "tolist"):
ids = ids.tolist()
ids = to_native_list(ids)

# Handle array values with bulk_update_mappings
schemas = self._attr_schemas_for_table(table_class)
Expand Down Expand Up @@ -2255,8 +2248,8 @@ def _get_degree(
) -> list[int] | int:
edge_key_col = getattr(self.Edge, node_key)

if isinstance(node_ids, int):
stmt = sa.select(sa.func.count()).where(edge_key_col == node_ids)
if is_int_like(node_ids):
stmt = sa.select(sa.func.count()).where(edge_key_col == to_native(node_ids))
with Session(self._engine) as session:
return int(session.execute(stmt).scalar())

Expand All @@ -2267,6 +2260,7 @@ def _get_degree(
if node_ids is None:
degree.update(session.execute(base_stmt).all())
else:
node_ids = to_native_list(node_ids)
# Chunk the IN(...) so the bound-parameter count stays below
# the backend's limit (notably SQLite's
# ``SQLITE_MAX_VARIABLE_NUMBER``). Each chunk's group-by result
Expand Down Expand Up @@ -2590,7 +2584,7 @@ def has_node(self, node_id: int) -> bool:
Check if the graph has a node with the given id.
"""
with Session(self._engine) as session:
return session.scalar(sa.sql.expression.exists().where(self.Node.node_id == node_id).select())
return session.scalar(sa.sql.expression.exists().where(self.Node.node_id == to_native(node_id)).select())

def has_edge(self, source_id: int, target_id: int) -> bool:
"""
Expand All @@ -2599,7 +2593,7 @@ def has_edge(self, source_id: int, target_id: int) -> bool:
with Session(self._engine) as session:
return (
session.query(self.Edge)
.filter(self.Edge.source_id == source_id, self.Edge.target_id == target_id)
.filter(self.Edge.source_id == to_native(source_id), self.Edge.target_id == to_native(target_id))
.count()
> 0
)
Expand All @@ -2611,7 +2605,7 @@ def edge_id(self, source_id: int, target_id: int) -> int:
with Session(self._engine) as session:
edge_id = (
session.query(self.Edge.edge_id)
.filter(self.Edge.source_id == source_id, self.Edge.target_id == target_id)
.filter(self.Edge.source_id == to_native(source_id), self.Edge.target_id == to_native(target_id))
.scalar()
)
if edge_id is None:
Expand All @@ -2632,10 +2626,7 @@ def bulk_remove_edges(self, edge_ids: Sequence[int]) -> None:
ValueError
If any edge_id does not exist in the graph.
"""
if hasattr(edge_ids, "tolist"):
edge_ids = edge_ids.tolist()
else:
edge_ids = list(edge_ids)
edge_ids = to_native_list(edge_ids)

if len(edge_ids) == 0:
return
Expand Down
Loading
Loading