diff --git a/src/funlib/persistence/graphs/sql_graph_database.py b/src/funlib/persistence/graphs/sql_graph_database.py index 5cb511d..e1c0811 100644 --- a/src/funlib/persistence/graphs/sql_graph_database.py +++ b/src/funlib/persistence/graphs/sql_graph_database.py @@ -2,6 +2,7 @@ from abc import abstractmethod from typing import Any, Iterable, Optional +import numpy as np from funlib.geometry import Coordinate, Roi from networkx import DiGraph, Graph from networkx.classes.reportviews import NodeView, OutEdgeView @@ -704,10 +705,22 @@ def update_nodes( ) continue - values = [data.get(attr) for attr in attrs] - setters = [ - f"{k} = {self.__convert_to_sql(v)}" for k, v in zip(attrs, values) - ] + setters = [] + for attr in attrs: + value = data.get(attr) + columns = self._node_attrs_to_columns([attr]) + if len(columns) == 1: + setters.append(f"{columns[0]} = {self.__convert_to_sql(value)}") + else: + # attr is stored across multiple columns (e.g. a Vec in the + # SQLite backend -> attr_0, attr_1, ...); expand the value to + # one assignment per column instead of writing the whole + # sequence into a single column (which produced invalid SQL). + components = [None] * len(columns) if value is None else list(value) + setters.extend( + f"{col} = {self.__convert_to_sql(component)}" + for col, component in zip(columns, components) + ) update_statement = ( f"UPDATE {self.nodes_table_name} SET " f"{', '.join(setters)} WHERE " @@ -806,12 +819,22 @@ def __get_node_pos(self, n: dict[str, Any]) -> Optional[Coordinate]: return None def __convert_to_sql(self, x: Any) -> str: + # normalize numpy values to their native python equivalents; all + # other types pass through untouched + if isinstance(x, np.ndarray): + x = x.tolist() + elif isinstance(x, np.generic): + x = x.item() if isinstance(x, str): return f"'{x}'" elif x is None: return "null" elif isinstance(x, bool): return f"{x}".lower() + elif isinstance(x, (list, tuple)): + # a native array column (e.g. a Vec stored as a single array column + # in the PostgreSQL backend) -> array literal '{v0,v1,...}' + return "'{" + ",".join(map(str, x)) + "}'" else: return str(x) diff --git a/tests/test_graph.py b/tests/test_graph.py index eef1e4e..858dd76 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -1,4 +1,5 @@ import networkx as nx +import numpy as np import pytest from funlib.geometry import Roi @@ -180,6 +181,44 @@ def test_graph_read_and_update_specific_attrs(provider_factory): assert data["c"] == 5 +def test_graph_update_vector_node_attr(provider_factory): + # Regression: update_nodes must handle a Vec (multi-column) node attribute, + # not just scalars. Previously the SQLite backend emitted + # `UPDATE nodes SET embedding = [1.0, 2.0, 3.0]` (invalid SQL) because it + # wrote the whole sequence into one column instead of expanding to the + # per-dimension columns (embedding_0, embedding_1, ...). + roi = Roi((0, 0, 0), (10, 10, 10)) + graph = nx.Graph() + graph.add_node(7, position=(5, 5, 5), embedding=[0.0, 0.0, 0.0], score=0.0) + + with provider_factory( + "w", + node_attrs={ + "position": Vec(float, 3), + "embedding": Vec(float, 3), + "score": float, + }, + ) as gp: + gp.write_graph(graph) + + with provider_factory("r+") as gp: + g = gp.read_graph(roi) + # numpy values (the typical case, e.g. embeddings computed with numpy) + # must be normalized to native python on their way into the SQL + # statement, for arrays and scalars alike + g.nodes[7]["embedding"] = np.array([1.0, 2.0, 3.0]) + g.nodes[7]["score"] = np.float64(0.5) + try: + gp.update_nodes(g.nodes, roi=roi, attributes=["embedding", "score"]) + except NotImplementedError: + pytest.xfail() + + with provider_factory("r") as gp: + updated = gp.read_graph(roi) + assert list(updated.nodes[7]["embedding"]) == [1.0, 2.0, 3.0] + assert updated.nodes[7]["score"] == 0.5 + + def test_graph_read_unbounded_roi(provider_factory, write_method): unbounded_roi = Roi((None, None, None), (None, None, None)) graph = nx.Graph()