From 8f6fd987365c9d4c02b0906141aa5139f1a555b5 Mon Sep 17 00:00:00 2001 From: Teun Huijben Date: Thu, 30 Jul 2026 10:31:47 -0700 Subject: [PATCH 1/6] polars array reading inconsistent dtypes --- src/tracksdata/graph/_sql_graph.py | 49 +++++++++++++------ .../graph/_test/test_graph_backends.py | 23 +++++++++ src/tracksdata/utils/_dataframe.py | 28 +++++++++-- 3 files changed, 82 insertions(+), 18 deletions(-) diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index b2f66b46..541c16bd 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/src/tracksdata/graph/_sql_graph.py @@ -422,7 +422,7 @@ def _read_attr_dataframe(self, query: sa.Select, table: type[DeclarativeBase]) - schema_overrides=self._graph._polars_schema_override(table), ) - df = unpickle_bytes_columns(df) + df = unpickle_bytes_columns(df, self._graph._pickled_column_dtypes(table)) return self._graph._cast_columns(table, df) def _query_from_attr_keys( @@ -845,27 +845,46 @@ def _restore_pickled_column_types(self, table: sa.Table) -> None: if isinstance(column.type, sa.LargeBinary): column.type = sa.PickleType() - def _polars_schema_override(self, table_class: type[DeclarativeBase]) -> SchemaDict: - """Return polars dtype overrides for physical columns in *table_class*. + def _declared_column_dtypes(self, table_class: type[DeclarativeBase], *, pickled: bool) -> SchemaDict: + """Return the declared polars dtype of the physical columns in *table_class*. Flat struct leaf columns are included with their native leaf dtypes. - Pickled columns are excluded here and handled in a second pass by - ``_cast_array_columns``. + + Parameters + ---------- + table_class : type[DeclarativeBase] + The table to describe. + pickled : bool + Whether to return the columns stored as pickled blobs (arrays, lists, + objects, ...) or the ones stored as native SQL scalars. """ - overrides: SchemaDict = {} + dtypes: SchemaDict = {} schemas = self._attr_schemas_for_table(table_class) table_cols = table_class.__table__.columns for key, schema in schemas.items(): if isinstance(schema.dtype, pl.Struct): - # Emit overrides for each leaf physical column. + # Emit one entry per leaf physical column. for flat_col, leaf_dtype in flatten_struct_dtype(key, schema.dtype): - if flat_col in table_cols and not self._is_pickled_sql_type(table_cols[flat_col].type): - overrides[flat_col] = leaf_dtype - elif key in table_cols and not self._is_pickled_sql_type(table_cols[key].type): - overrides[key] = schema.dtype + if flat_col in table_cols and self._is_pickled_sql_type(table_cols[flat_col].type) == pickled: + dtypes[flat_col] = leaf_dtype + elif key in table_cols and self._is_pickled_sql_type(table_cols[key].type) == pickled: + dtypes[key] = schema.dtype + + return dtypes + + def _polars_schema_override(self, table_class: type[DeclarativeBase]) -> SchemaDict: + """Return polars dtype overrides for the natively stored columns in *table_class*. + + Pickled columns are excluded here: the raw query returns their bytes, so + their dtype can only be applied once the blobs have been unpickled, which + ``unpickle_bytes_columns`` does with ``_pickled_column_dtypes``. + """ + return self._declared_column_dtypes(table_class, pickled=False) - return overrides + def _pickled_column_dtypes(self, table_class: type[DeclarativeBase]) -> SchemaDict: + """Return the declared polars dtype of the pickled columns in *table_class*.""" + return self._declared_column_dtypes(table_class, pickled=True) @staticmethod def _build_struct_expr(key: str, dtype: pl.Struct) -> pl.Expr: @@ -1357,7 +1376,7 @@ def _get_neighbors( filter_node_ids, self.Node, ) - node_df = unpickle_bytes_columns(node_df) + node_df = unpickle_bytes_columns(node_df, self._pickled_column_dtypes(self.Node)) node_df = self._cast_columns(self.Node, node_df) if single_node: @@ -1550,7 +1569,7 @@ def node_attrs( connection=session.connection(), schema_overrides=self._polars_schema_override(self.Node), ) - nodes_df = unpickle_bytes_columns(nodes_df) + nodes_df = unpickle_bytes_columns(nodes_df, self._pickled_column_dtypes(self.Node)) nodes_df = self._cast_columns(self.Node, nodes_df) # Select using logical keys (struct columns are now reconstructed). @@ -1596,7 +1615,7 @@ def edge_attrs( connection=session.connection(), schema_overrides=self._polars_schema_override(self.Edge), ) - edges_df = unpickle_bytes_columns(edges_df) + edges_df = unpickle_bytes_columns(edges_df, self._pickled_column_dtypes(self.Edge)) edges_df = self._cast_columns(self.Edge, edges_df) if unpack: diff --git a/src/tracksdata/graph/_test/test_graph_backends.py b/src/tracksdata/graph/_test/test_graph_backends.py index 3c8095c6..7cea3e6a 100644 --- a/src/tracksdata/graph/_test/test_graph_backends.py +++ b/src/tracksdata/graph/_test/test_graph_backends.py @@ -109,6 +109,29 @@ def test_add_edge(graph_backend: BaseGraph) -> None: assert df["weight"].to_list() == [0.5, 0.1] +def test_array_attr_read_honors_declared_dtype(graph_backend: BaseGraph) -> None: + """An `Array(Float64)` column must not be truncated to integers when read back. + + The declared dtype has to win over any dtype inferred from the leading rows, + otherwise a whole-numbered first row silently truncates the fractional ones. + """ + graph_backend.add_node_attr_key("pos", dtype=pl.Array(pl.Float64, 2)) + graph_backend.add_node_attr_key("values", dtype=pl.List(pl.Float64)) + + graph_backend.bulk_add_nodes( + [ + {"t": 0, "pos": [50, 50], "values": [50, 50]}, # whole numbers + {"t": 1, "pos": [1.5, 1.5], "values": [1.5, 1.5]}, # fractional + ] + ) + + nodes_df = graph_backend.node_attrs(attr_keys=["t", "pos", "values"]).sort("t") + assert nodes_df.schema["pos"] == pl.Array(pl.Float64, 2) + assert nodes_df.schema["values"] == pl.List(pl.Float64) + assert nodes_df["pos"].to_list() == [[50.0, 50.0], [1.5, 1.5]] + assert nodes_df["values"].to_list() == [[50.0, 50.0], [1.5, 1.5]] + + def test_remove_edge_by_id(graph_backend: BaseGraph) -> None: """Test removing an edge by ID across backends using unified API.""" # Setup diff --git a/src/tracksdata/utils/_dataframe.py b/src/tracksdata/utils/_dataframe.py index a6de0f17..8c692081 100644 --- a/src/tracksdata/utils/_dataframe.py +++ b/src/tracksdata/utils/_dataframe.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping + import cloudpickle import polars as pl import polars.selectors as cs @@ -29,7 +31,10 @@ def unpack_array_attrs(df: pl.DataFrame) -> pl.DataFrame: return unpack_array_attrs(df) -def unpickle_bytes_columns(df: pl.DataFrame) -> pl.DataFrame: +def unpickle_bytes_columns( + df: pl.DataFrame, + dtypes: Mapping[str, pl.DataType] | None = None, +) -> pl.DataFrame: """ Unpickle bytes columns from the database. @@ -37,17 +42,34 @@ def unpickle_bytes_columns(df: pl.DataFrame) -> pl.DataFrame: ---------- df : pl.DataFrame The DataFrame to unpickle the bytes columns from. + dtypes : Mapping[str, pl.DataType] | None + Declared dtype per column, used to build the unpickled columns. + Columns without a declared dtype fall back to polars' inference, which + only looks at the leading rows and therefore silently truncates a + `Float64` column whose first rows happen to hold whole numbers. Returns ------- pl.DataFrame The DataFrame with the bytes columns unpickled. """ + if dtypes is None: + dtypes = {} + 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): + if not isinstance(dtype, pl.Object): + continue + values = df[col].to_list() + # `None` falls back to polars' inference, either because the column has no + # declared dtype or because its values turned out not to fit it. + candidates = (dtypes[col], None) if col in dtypes else (None,) + for target_dtype in candidates: try: - df = df.with_columns(pl.Series(df[col].to_list()).alias(col)) + df = df.with_columns(pl.Series(col, values, dtype=target_dtype)) + break except Exception: + # values that fit neither the declared dtype nor an inferred one + # (e.g. `Mask` objects) are left as an object column. pass return df From 1640a8b816a858c396c624ef07b115b434740dd6 Mon Sep 17 00:00:00 2001 From: Jordao Bragantini Date: Fri, 7 Aug 2026 13:44:43 -0700 Subject: [PATCH 2/6] Refactor SQLGraph database reads --- src/tracksdata/graph/_sql_graph.py | 97 ++++++++++++------------------ 1 file changed, 40 insertions(+), 57 deletions(-) diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index 541c16bd..2e58da2b 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/src/tracksdata/graph/_sql_graph.py @@ -403,7 +403,7 @@ def node_attrs( attr_keys=attr_keys, ) - nodes_attrs = self._read_attr_dataframe(query, self._graph.Node) + nodes_attrs = self._graph._read_database(query, self._graph.Node) if attr_keys is not None: attr_keys = list(dict.fromkeys(attr_keys)) @@ -414,17 +414,6 @@ def node_attrs( return nodes_attrs - def _read_attr_dataframe(self, query: sa.Select, table: type[DeclarativeBase]) -> pl.DataFrame: - with Session(self._graph._engine) as session: - df = pl.read_database( - self._graph._raw_query(query), - connection=session.connection(), - schema_overrides=self._graph._polars_schema_override(table), - ) - - df = unpickle_bytes_columns(df, self._graph._pickled_column_dtypes(table)) - return self._graph._cast_columns(table, df) - def _query_from_attr_keys( self, query: sa.Select, @@ -469,7 +458,7 @@ def edge_attrs(self, attr_keys: list[str] | None = None, unpack: bool = False) - ], ) - edges_df = self._read_attr_dataframe(query, self._graph.Edge) + edges_df = self._graph._read_database(query, self._graph.Edge) if unpack: edges_df = unpack_array_attrs(edges_df) @@ -514,8 +503,8 @@ def subgraph( ], ) - nodes_df = self._read_attr_dataframe(node_query, self._graph.Node) - edges_df = self._read_attr_dataframe(edge_query, self._graph.Edge) + nodes_df = self._graph._read_database(node_query, self._graph.Node) + edges_df = self._graph._read_database(edge_query, self._graph.Edge) node_map_to_root = {} node_map_from_root = {} @@ -862,29 +851,46 @@ def _declared_column_dtypes(self, table_class: type[DeclarativeBase], *, pickled schemas = self._attr_schemas_for_table(table_class) table_cols = table_class.__table__.columns + flat_key_to_dtype: list[tuple[str, pl.DataType]] = [] + # flatten structs into their leaf columns for key, schema in schemas.items(): if isinstance(schema.dtype, pl.Struct): - # Emit one entry per leaf physical column. - for flat_col, leaf_dtype in flatten_struct_dtype(key, schema.dtype): - if flat_col in table_cols and self._is_pickled_sql_type(table_cols[flat_col].type) == pickled: - dtypes[flat_col] = leaf_dtype - elif key in table_cols and self._is_pickled_sql_type(table_cols[key].type) == pickled: - dtypes[key] = schema.dtype + flat_key_to_dtype.extend( + (flat_col, leaf_dtype) for flat_col, leaf_dtype in flatten_struct_dtype(key, schema.dtype) + ) + else: + flat_key_to_dtype.append((key, schema.dtype)) + + for key, dtype in flat_key_to_dtype: + if key in table_cols and self._is_pickled_sql_type(table_cols[key].type) == pickled: + dtypes[key] = dtype return dtypes - def _polars_schema_override(self, table_class: type[DeclarativeBase]) -> SchemaDict: - """Return polars dtype overrides for the natively stored columns in *table_class*. + def _read_database( + self, + query: sa.Select, + table_class: type[DeclarativeBase], + connection: sa.Connection | None = None, + ) -> pl.DataFrame: + """Read a SQL query and restore the declared Polars attribute dtypes. - Pickled columns are excluded here: the raw query returns their bytes, so - their dtype can only be applied once the blobs have been unpickled, which - ``unpickle_bytes_columns`` does with ``_pickled_column_dtypes``. + Native SQL columns receive schema overrides during the database read. + Pickled columns are unpickled before their declared dtypes are restored, + and flat struct columns are reconstructed into logical struct columns. + A temporary session supplies the connection when one is not provided. """ - return self._declared_column_dtypes(table_class, pickled=False) + if connection is None: + with Session(self._engine) as session: + return self._read_database(query, table_class, session.connection()) - def _pickled_column_dtypes(self, table_class: type[DeclarativeBase]) -> SchemaDict: - """Return the declared polars dtype of the pickled columns in *table_class*.""" - return self._declared_column_dtypes(table_class, pickled=True) + df = pl.read_database( + self._raw_query(query), + connection=connection, + schema_overrides=self._declared_column_dtypes(table_class, pickled=False), + ) + df = unpickle_bytes_columns(df, self._declared_column_dtypes(table_class, pickled=True)) + return self._cast_columns(table_class, df) @staticmethod def _build_struct_expr(key: str, dtype: pl.Struct) -> pl.Expr: @@ -1364,11 +1370,7 @@ def _get_neighbors( query = session.query(getattr(self.Edge, node_key), *node_columns) query = query.join(self.Edge, getattr(self.Edge, neighbor_key) == self.Node.node_id) if filter_node_ids is None or len(filter_node_ids) == 0: - node_df = pl.read_database( - query.statement, - connection=session.connection(), - schema_overrides=self._polars_schema_override(self.Node), - ) + node_df = self._read_database(query.statement, self.Node, session.connection()) else: node_df = self._chunked_sa_read( session, @@ -1376,8 +1378,6 @@ def _get_neighbors( filter_node_ids, self.Node, ) - node_df = unpickle_bytes_columns(node_df, self._pickled_column_dtypes(self.Node)) - node_df = self._cast_columns(self.Node, node_df) if single_node: if not return_attrs: @@ -1564,13 +1564,7 @@ def node_attrs( *self._physical_cols_for_query(attr_keys, self.Node), ) - nodes_df = pl.read_database( - self._raw_query(query), - connection=session.connection(), - schema_overrides=self._polars_schema_override(self.Node), - ) - nodes_df = unpickle_bytes_columns(nodes_df, self._pickled_column_dtypes(self.Node)) - nodes_df = self._cast_columns(self.Node, nodes_df) + nodes_df = self._read_database(query, self.Node, session.connection()) # Select using logical keys (struct columns are now reconstructed). if attr_keys is not None: @@ -1610,13 +1604,7 @@ def edge_attrs( *self._physical_cols_for_query(attr_keys, self.Edge), ) - edges_df = pl.read_database( - self._raw_query(query), - connection=session.connection(), - schema_overrides=self._polars_schema_override(self.Edge), - ) - edges_df = unpickle_bytes_columns(edges_df, self._pickled_column_dtypes(self.Edge)) - edges_df = self._cast_columns(self.Edge, edges_df) + edges_df = self._read_database(query, self.Edge, session.connection()) if unpack: edges_df = unpack_array_attrs(edges_df) @@ -2151,12 +2139,7 @@ def _chunked_sa_read( chunks = [] for i in range(0, len(data), chunk_size): query = query_filter_op(data[i : i + chunk_size]) - data_df = pl.read_database( - query.statement, - connection=session.connection(), - schema_overrides=self._polars_schema_override(table_class), - ) - chunks.append(data_df) + chunks.append(self._read_database(query.statement, table_class, session.connection())) return pl.concat(chunks) def _create_id_scratch_table(self, ids: Sequence[int]) -> sa.Table: From 848607c2133dbf5290eef6fa63a3c03724e7b08c Mon Sep 17 00:00:00 2001 From: Jordao Bragantini Date: Fri, 7 Aug 2026 13:55:03 -0700 Subject: [PATCH 3/6] Simplify SQLGraph dataframe reads --- src/tracksdata/graph/_sql_graph.py | 147 +++++++++++------------------ 1 file changed, 55 insertions(+), 92 deletions(-) diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index 2e58da2b..151fb3e7 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/src/tracksdata/graph/_sql_graph.py @@ -834,39 +834,6 @@ def _restore_pickled_column_types(self, table: sa.Table) -> None: if isinstance(column.type, sa.LargeBinary): column.type = sa.PickleType() - def _declared_column_dtypes(self, table_class: type[DeclarativeBase], *, pickled: bool) -> SchemaDict: - """Return the declared polars dtype of the physical columns in *table_class*. - - Flat struct leaf columns are included with their native leaf dtypes. - - Parameters - ---------- - table_class : type[DeclarativeBase] - The table to describe. - pickled : bool - Whether to return the columns stored as pickled blobs (arrays, lists, - objects, ...) or the ones stored as native SQL scalars. - """ - dtypes: SchemaDict = {} - schemas = self._attr_schemas_for_table(table_class) - table_cols = table_class.__table__.columns - - flat_key_to_dtype: list[tuple[str, pl.DataType]] = [] - # flatten structs into their leaf columns - for key, schema in schemas.items(): - if isinstance(schema.dtype, pl.Struct): - flat_key_to_dtype.extend( - (flat_col, leaf_dtype) for flat_col, leaf_dtype in flatten_struct_dtype(key, schema.dtype) - ) - else: - flat_key_to_dtype.append((key, schema.dtype)) - - for key, dtype in flat_key_to_dtype: - if key in table_cols and self._is_pickled_sql_type(table_cols[key].type) == pickled: - dtypes[key] = dtype - - return dtypes - def _read_database( self, query: sa.Select, @@ -884,81 +851,77 @@ def _read_database( with Session(self._engine) as session: return self._read_database(query, table_class, session.connection()) + native_dtypes, pickled_dtypes, struct_dtypes = self._database_column_dtypes(table_class) df = pl.read_database( self._raw_query(query), connection=connection, - schema_overrides=self._declared_column_dtypes(table_class, pickled=False), + schema_overrides=native_dtypes, ) - df = unpickle_bytes_columns(df, self._declared_column_dtypes(table_class, pickled=True)) - return self._cast_columns(table_class, df) - - @staticmethod - def _build_struct_expr(key: str, dtype: pl.Struct) -> pl.Expr: - """Recursively build a ``pl.struct`` expression from flat leaf columns.""" - fields: list[pl.Expr] = [] - for field_name, field_dtype in dtype.to_schema().items(): - flat_col = f"{key}{STRUCT_FIELD_SEP}{field_name}" - if isinstance(field_dtype, pl.Struct): - fields.append(SQLGraph._build_struct_expr(flat_col, field_dtype).alias(field_name)) - else: - fields.append(pl.col(flat_col).alias(field_name)) - return pl.struct(fields) + df = unpickle_bytes_columns(df, pickled_dtypes) + return self._reconstruct_struct_columns(df, struct_dtypes) - def _cast_columns(self, table_class: type[DeclarativeBase], df: pl.DataFrame) -> pl.DataFrame: - """Cast pickled columns to their target dtype and reconstruct struct columns.""" - schemas = self._attr_schemas_for_table(table_class) + def _database_column_dtypes( + self, + table_class: type[DeclarativeBase], + ) -> tuple[SchemaDict, SchemaDict, dict[str, pl.Struct]]: + """Partition physical column dtypes by storage and collect logical structs.""" + native_dtypes: SchemaDict = {} + pickled_dtypes: SchemaDict = {} + struct_dtypes: dict[str, pl.Struct] = {} table_cols = table_class.__table__.columns - casts: list[pl.Series] = [] - struct_keys: list[tuple[str, pl.Struct]] = [] + for key, schema in self._attr_schemas_for_table(table_class).items(): + is_struct = isinstance(schema.dtype, pl.Struct) + if is_struct: + struct_dtypes[key] = schema.dtype + physical_dtypes = flatten_struct_dtype(key, schema.dtype) if is_struct else ((key, schema.dtype),) - for key, schema in schemas.items(): - if isinstance(schema.dtype, pl.Struct): - # Cast any pickled flat leaf columns to their proper dtypes before - # reconstruction so Array/List fields have correct dtype. - for flat_col, leaf_dtype in flatten_struct_dtype(key, schema.dtype): - if flat_col not in df.columns or flat_col not in table_cols: - continue - if not self._is_pickled_sql_type(table_cols[flat_col].type): - continue - try: - casts.append(pl.Series(flat_col, df[flat_col].to_list(), dtype=leaf_dtype)) - except Exception: - continue - struct_keys.append((key, schema.dtype)) - continue + for column_name, dtype in physical_dtypes: + if column_name not in table_cols: + continue + target = pickled_dtypes if self._is_pickled_sql_type(table_cols[column_name].type) else native_dtypes + target[column_name] = dtype - if key not in df.columns or key not in table_cols: - continue + return native_dtypes, pickled_dtypes, struct_dtypes - if not self._is_pickled_sql_type(table_cols[key].type): - continue - - try: - casts.append(pl.Series(key, df[key].to_list(), dtype=schema.dtype)) - except Exception: - # Keep original dtype when values cannot be cast to the target schema. + def _reconstruct_struct_columns( + self, + df: pl.DataFrame, + struct_dtypes: dict[str, pl.Struct], + ) -> pl.DataFrame: + """Reconstruct logical struct columns from flat physical columns.""" + struct_exprs: list[pl.Expr] = [] + flat_cols_to_drop: list[str] = [] + for key, dtype in struct_dtypes.items(): + flat_cols = [column_name for column_name, _ in flatten_struct_dtype(key, dtype)] + missing_cols = [column_name for column_name in flat_cols if column_name not in df.columns] + if len(missing_cols) == len(flat_cols): continue - - if casts: - df = df.with_columns(casts) - - # Reconstruct struct columns from their flat physical columns. - for key, dtype in struct_keys: - flat_cols = [fc for fc, _ in flatten_struct_dtype(key, dtype)] - present = [fc for fc in flat_cols if fc in df.columns] - if not present: - continue # struct was not part of this query; skip - missing = [fc for fc in flat_cols if fc not in df.columns] - if missing: + if missing_cols: raise ValueError( f"Struct attribute '{key}' is partially present in the DataFrame " - f"(missing: {missing}). Cannot reconstruct the struct column." + f"(missing: {missing_cols}). Cannot reconstruct the struct column." ) - df = df.with_columns(self._build_struct_expr(key, dtype).alias(key)).drop(flat_cols) + struct_exprs.append(self._build_struct_expr(key, dtype).alias(key)) + flat_cols_to_drop.extend(flat_cols) + + if struct_exprs: + df = df.with_columns(struct_exprs).drop(flat_cols_to_drop) return df + @staticmethod + def _build_struct_expr(key: str, dtype: pl.Struct) -> pl.Expr: + """Recursively build a ``pl.struct`` expression from flat leaf columns.""" + fields: list[pl.Expr] = [] + for field_name, field_dtype in dtype.to_schema().items(): + flat_col = f"{key}{STRUCT_FIELD_SEP}{field_name}" + if isinstance(field_dtype, pl.Struct): + fields.append(SQLGraph._build_struct_expr(flat_col, field_dtype).alias(field_name)) + else: + fields.append(pl.col(flat_col).alias(field_name)) + return pl.struct(fields) + def _update_max_id_per_time(self) -> None: """ Update the maximum node ID for each time point. @@ -1640,7 +1603,7 @@ def _physical_column_names( Logical keys are what the user sees (``"measurements"``); physical columns are what actually exists in the table (``"measurements__score"``, ...). The two - diverge only for struct attributes; ``_cast_columns`` reassembles the struct + diverge only for struct attributes; ``_read_database`` reassembles the struct on the result DataFrame. """ schemas = self._attr_schemas_for_table(table_class) From 579ca120846061c23abdec3dbdce468c7eff9668 Mon Sep 17 00:00:00 2001 From: Teun Huijben <45037215+TeunHuijben@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:08:21 -0700 Subject: [PATCH 4/6] fix metadata loss in to_geff (#333) * fix metadata loss in to_geff * merge two io/geff files into 1 --- src/tracksdata/graph/_base_graph.py | 25 +++++- .../graph/_test/test_graph_backends.py | 50 +++++++++++ src/tracksdata/io/__init__.py | 7 +- .../io/{_geff_dtypes.py => _geff.py} | 76 +++++++++++++--- src/tracksdata/io/_test/test_geff_dtypes.py | 2 +- src/tracksdata/io/_test/test_geff_metadata.py | 89 +++++++++++++++++++ 6 files changed, 229 insertions(+), 20 deletions(-) rename src/tracksdata/io/{_geff_dtypes.py => _geff.py} (72%) create mode 100644 src/tracksdata/io/_test/test_geff_metadata.py diff --git a/src/tracksdata/graph/_base_graph.py b/src/tracksdata/graph/_base_graph.py index 3a214e35..51f7bf6f 100644 --- a/src/tracksdata/graph/_base_graph.py +++ b/src/tracksdata/graph/_base_graph.py @@ -1868,6 +1868,14 @@ def to_geff( It automatically generates the metadata with: - axes: time (t) and spatial axes ((z), y, x) - tracklet node property: tracklet_id + The graph metadata (`graph.metadata`) is always written to + `geff_metadata.extra["tracksdata"]`, including when the metadata is provided + by the caller. On key collisions the caller's value wins, so an explicit + `extra["tracksdata"]` entry still overrides the graph's metadata. + The caller's object is not modified. + `shape` is the canonical key for the shape of the dense segmentation, it is + read back by `GraphArrayView` and `to_ctc`. Use + `tracksdata.io.read_graph_metadata` to read it back without building a graph. overwrite : bool Whether to overwrite the geff data directory if it exists. zarr_format : Literal[2, 3] @@ -1883,6 +1891,10 @@ def to_geff( edge_ids = edge_attrs.select(DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET).to_numpy() edge_attrs = edge_attrs.drop(DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET) + td_metadata = self.metadata.copy() + td_metadata.update(self._private_metadata_for_copy()) + td_metadata.pop("geff", None) # avoid geff being written multiple times + if geff_metadata is None: axes = [Axis(name=DEFAULT_ATTR_KEYS.T, type="time")] axes.extend( @@ -1917,10 +1929,6 @@ def to_geff( for k, v in edge_attrs.to_dict().items() } - td_metadata = self.metadata.copy() - td_metadata.update(self._private_metadata_for_copy()) - td_metadata.pop("geff", None) # avoid geff being written multiple times - geff_metadata = geff.GeffMetadata( directed=True, axes=axes, @@ -1931,6 +1939,15 @@ def to_geff( "tracksdata": td_metadata, }, ) + else: + # copy so the caller's metadata object is left untouched + geff_metadata = geff_metadata.model_copy(deep=True) + extra = dict(geff_metadata.extra) + # caller-provided entries win, so they can still override the graph's metadata + merged = {**td_metadata, **extra.get("tracksdata", {})} + merged.pop("geff", None) # avoid geff being written multiple times + extra["tracksdata"] = merged + geff_metadata.extra = extra node_dict = { k: {"values": column_to_numpy(v), "missing": None} diff --git a/src/tracksdata/graph/_test/test_graph_backends.py b/src/tracksdata/graph/_test/test_graph_backends.py index 7cea3e6a..15cb0a99 100644 --- a/src/tracksdata/graph/_test/test_graph_backends.py +++ b/src/tracksdata/graph/_test/test_graph_backends.py @@ -9,6 +9,7 @@ import pytest import rustworkx as rx import sqlalchemy as sa +from geff_spec import GeffMetadata from zarr.storage import MemoryStore from tracksdata.attrs import EdgeAttr, NodeAttr @@ -2884,6 +2885,55 @@ def test_geff_roundtrip(graph_backend: BaseGraph) -> None: ) +def test_geff_roundtrip_custom_metadata(graph_backend: BaseGraph) -> None: + """Graph metadata must survive `to_geff` when the caller supplies its own `GeffMetadata`.""" + + _fill_mock_geff_graph(graph_backend) + graph_backend.metadata["shape"] = (5, 100, 100) + + # a downstream library supplying its own metadata: same props as the auto-generated + # one, but with its own `extra` namespace instead of tracksdata's. + reference_store = MemoryStore() + graph_backend.to_geff(geff_store=reference_store) + custom_metadata = GeffMetadata.read(reference_store) + custom_metadata.extra = {"downstream": {"hello": "world"}} + + output_store = MemoryStore() + graph_backend.to_geff(geff_store=output_store, geff_metadata=custom_metadata) + + written_metadata = GeffMetadata.read(output_store) + # the caller's own namespace is untouched ... + assert written_metadata.extra["downstream"] == {"hello": "world"} + # ... and the graph metadata rode along. `shape` is a list, not a tuple, because + # the extras are serialized as JSON. + assert written_metadata.extra["tracksdata"]["shape"] == [5, 100, 100] + + geff_graph, _ = IndexedRXGraph.from_geff(output_store) + assert geff_graph.metadata["shape"] == [5, 100, 100] + + # the metadata object the caller passed in was not modified + assert custom_metadata.extra == {"downstream": {"hello": "world"}} + + +def test_geff_custom_metadata_overrides_graph_metadata(graph_backend: BaseGraph) -> None: + """On key collisions the caller-supplied `extra["tracksdata"]` wins.""" + + _fill_mock_geff_graph(graph_backend) + graph_backend.metadata["shape"] = (5, 100, 100) + + reference_store = MemoryStore() + graph_backend.to_geff(geff_store=reference_store) + custom_metadata = GeffMetadata.read(reference_store) + custom_metadata.extra = {"tracksdata": {"shape": [1, 2, 3], "extra_key": "value"}} + + output_store = MemoryStore() + graph_backend.to_geff(geff_store=output_store, geff_metadata=custom_metadata) + + geff_graph, _ = IndexedRXGraph.from_geff(output_store) + assert geff_graph.metadata["shape"] == [1, 2, 3] + assert geff_graph.metadata["extra_key"] == "value" + + def test_geff_overwrite(graph_backend: BaseGraph, tmp_path: Path) -> None: """Test that to_geff overwrites existing data in the store.""" _fill_mock_geff_graph(graph_backend) diff --git a/src/tracksdata/io/__init__.py b/src/tracksdata/io/__init__.py index 57313e1a..e502138a 100644 --- a/src/tracksdata/io/__init__.py +++ b/src/tracksdata/io/__init__.py @@ -1,12 +1,17 @@ """Input/output utilities for loading and saving tracking data in various formats.""" from tracksdata.io._ctc import compressed_tracks_table, from_ctc, to_ctc -from tracksdata.io._geff_dtypes import convert_geff_prop_dtype, geff_prop_dtype +from tracksdata.io._geff import ( + convert_geff_prop_dtype, + geff_prop_dtype, + read_graph_metadata, +) __all__ = [ "compressed_tracks_table", "convert_geff_prop_dtype", "from_ctc", "geff_prop_dtype", + "read_graph_metadata", "to_ctc", ] diff --git a/src/tracksdata/io/_geff_dtypes.py b/src/tracksdata/io/_geff.py similarity index 72% rename from src/tracksdata/io/_geff_dtypes.py rename to src/tracksdata/io/_geff.py index b29a4e24..ae954e16 100644 --- a/src/tracksdata/io/_geff_dtypes.py +++ b/src/tracksdata/io/_geff.py @@ -1,16 +1,24 @@ -"""Utilities for inspecting and converting the dtype of properties in geff files. - -The motivating case is segmentation masks: they are binary, but geff files -written by older versions of tracksdata stored the mask ``data`` buffer as -``uint64`` (see https://github.com/royerlab/tracksdata/pull/318). That is 8x -larger than a boolean buffer both on disk and, more importantly, when read into -memory, which can cause out-of-memory errors when loading large datasets. New -files store masks as ``bool`` at write time, so :func:`convert_geff_prop_dtype` -provides a one-time fix for legacy files. - -The helpers are not mask-specific: they read and rewrite the payload dtype of -any geff property (node or edge, fixed- or variable-length). The caller names -the property to act on. +"""Utilities for inspecting and repairing geff datasets without loading the graph. + +:func:`read_graph_metadata` reads the tracksdata graph metadata of a geff dataset. +`BaseGraph.to_geff` writes the graph metadata (`graph.metadata`) into the geff +metadata extras and `BaseGraph.from_geff` hoists it back onto the graph, but callers +that need a value *before* they have a graph object -- for example the ``shape`` of +the dense segmentation, which is required to construct a `GraphArrayView` -- can use +neither. This closes that gap so downstream libraries do not have to know where +tracksdata stores the extras. + +:func:`geff_prop_dtype` and :func:`convert_geff_prop_dtype` inspect and convert the +on-disk dtype of a property. The motivating case is segmentation masks: they are +binary, but geff files written by older versions of tracksdata stored the mask +``data`` buffer as ``uint64`` (see +https://github.com/royerlab/tracksdata/pull/318). That is 8x larger than a boolean +buffer both on disk and, more importantly, when read into memory, which can cause +out-of-memory errors when loading large datasets. New files store masks as ``bool`` +at write time, so :func:`convert_geff_prop_dtype` provides a one-time fix for legacy +files. The helpers are not mask-specific: they read and rewrite the payload dtype of +any geff property (node or edge, fixed- or variable-length). The caller names the +property to act on. """ from __future__ import annotations @@ -18,14 +26,54 @@ import os import shutil from pathlib import Path +from typing import Any import numpy as np import zarr +from geff_spec import GeffMetadata from zarr.storage import StoreLike +from tracksdata.graph._base_graph import BaseGraph from tracksdata.utils._logging import LOG -__all__ = ["convert_geff_prop_dtype", "geff_prop_dtype"] +__all__ = ["convert_geff_prop_dtype", "geff_prop_dtype", "read_graph_metadata"] + +_EXTRA_KEY = "tracksdata" + + +def read_graph_metadata(source: StoreLike | GeffMetadata) -> dict[str, Any]: + """ + Read the tracksdata graph metadata of a geff dataset without loading the graph. + + Returns the same metadata that `graph.metadata` would hold after + `BaseGraph.from_geff`, minus the `geff` key. Note that the values went through a + JSON round-trip, so tuples come back as lists. + + Parameters + ---------- + source : StoreLike | GeffMetadata + The store or path of the geff dataset, or an already parsed `GeffMetadata`. + + Returns + ------- + dict[str, Any] + The graph metadata, empty if the dataset was not written by tracksdata. + + Examples + -------- + ```python + graph.metadata["shape"] = (5, 100, 100) + graph.to_geff("tracks.geff") + + shape = read_graph_metadata("tracks.geff")["shape"] # [5, 100, 100] + ``` + """ + if not isinstance(source, GeffMetadata): + source = GeffMetadata.read(source) + + metadata = source.extra.get(_EXTRA_KEY, {}) + + return {k: v for k, v in metadata.items() if not BaseGraph._is_private_metadata_key(k)} def geff_prop_dtype( diff --git a/src/tracksdata/io/_test/test_geff_dtypes.py b/src/tracksdata/io/_test/test_geff_dtypes.py index 18733028..9f7354f1 100644 --- a/src/tracksdata/io/_test/test_geff_dtypes.py +++ b/src/tracksdata/io/_test/test_geff_dtypes.py @@ -8,7 +8,7 @@ from tracksdata.constants import DEFAULT_ATTR_KEYS from tracksdata.graph import IndexedRXGraph, RustWorkXGraph from tracksdata.io import convert_geff_prop_dtype, geff_prop_dtype -from tracksdata.io._geff_dtypes import _overwrite_array, _set_prop_metadata_dtype +from tracksdata.io._geff import _overwrite_array, _set_prop_metadata_dtype from tracksdata.nodes._mask import Mask MASK_KEY = DEFAULT_ATTR_KEYS.MASK diff --git a/src/tracksdata/io/_test/test_geff_metadata.py b/src/tracksdata/io/_test/test_geff_metadata.py new file mode 100644 index 00000000..2c577fc4 --- /dev/null +++ b/src/tracksdata/io/_test/test_geff_metadata.py @@ -0,0 +1,89 @@ +from pathlib import Path + +import polars as pl +from geff_spec import Axis, GeffMetadata, PropMetadata +from zarr.storage import MemoryStore + +from tracksdata.graph import RustWorkXGraph +from tracksdata.io import read_graph_metadata + +SHAPE = (5, 100, 100) + + +def _make_graph() -> RustWorkXGraph: + graph = RustWorkXGraph() + graph.add_node_attr_key("y", pl.Float64) + graph.add_node_attr_key("x", pl.Float64) + graph.add_node({"t": 0, "y": 1.0, "x": 2.0}) + graph.add_node({"t": 1, "y": 3.0, "x": 4.0}) + graph.metadata["shape"] = SHAPE + return graph + + +def _minimal_geff_metadata() -> GeffMetadata: + """A `GeffMetadata` as a downstream library would build it: no tracksdata extras.""" + return GeffMetadata( + directed=True, + axes=[ + Axis(name="t", type="time"), + Axis(name="y", type="space", scale=0.5), + Axis(name="x", type="space", scale=0.5), + ], + node_props_metadata={ + "t": PropMetadata(identifier="t", dtype="int64"), + "y": PropMetadata(identifier="y", dtype="float64"), + "x": PropMetadata(identifier="x", dtype="float64"), + }, + edge_props_metadata={}, + extra={"downstream": {"hello": "world"}}, + ) + + +def test_read_graph_metadata_from_path(tmp_path: Path) -> None: + """The shape is readable from a store path, without building a graph.""" + graph = _make_graph() + geff_path = tmp_path / "tracks.geff" + graph.to_geff(geff_store=geff_path) + + # tuples become lists through the JSON round-trip + assert read_graph_metadata(geff_path) == {"shape": list(SHAPE)} + + +def test_read_graph_metadata_custom_geff_metadata() -> None: + """The shape survives a write with caller-supplied metadata and is readable back.""" + graph = _make_graph() + store = MemoryStore() + graph.to_geff(geff_store=store, geff_metadata=_minimal_geff_metadata()) + + assert read_graph_metadata(store) == {"shape": list(SHAPE)} + + +def test_read_graph_metadata_from_geff_metadata_instance() -> None: + """An already parsed `GeffMetadata` is accepted, so the store is not reopened.""" + graph = _make_graph() + store = MemoryStore() + graph.to_geff(geff_store=store) + + assert read_graph_metadata(GeffMetadata.read(store)) == {"shape": list(SHAPE)} + + +def test_read_graph_metadata_without_tracksdata_extras() -> None: + """A geff not written by tracksdata yields an empty dict rather than raising.""" + # foreign extras only + assert read_graph_metadata(_minimal_geff_metadata()) == {} + + no_extra = _minimal_geff_metadata() + no_extra.extra = {} + assert read_graph_metadata(no_extra) == {} + + +def test_read_graph_metadata_excludes_private_keys() -> None: + """Private metadata is written to the store but not exposed by the reader.""" + graph = _make_graph() + graph._private_metadata["__private_secret"] = 42 + + store = MemoryStore() + graph.to_geff(geff_store=store) + + assert "__private_secret" in GeffMetadata.read(store).extra["tracksdata"] + assert read_graph_metadata(store) == {"shape": list(SHAPE)} From bbc51ab2a39bc43c0a5a179c12fe134b79c12c76 Mon Sep 17 00:00:00 2001 From: Teun Huijben <45037215+TeunHuijben@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:03:29 -0700 Subject: [PATCH 5/6] fix add_edge_to_view leaving edge_id to -1 in SQL (#331) --- src/tracksdata/graph/_graph_view.py | 5 +++++ src/tracksdata/graph/_test/test_subgraph.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/src/tracksdata/graph/_graph_view.py b/src/tracksdata/graph/_graph_view.py index 21d8a513..1be0cf55 100644 --- a/src/tracksdata/graph/_graph_view.py +++ b/src/tracksdata/graph/_graph_view.py @@ -798,6 +798,11 @@ def _add_edge_local(self, source_id: int, target_id: int) -> int: if c in df.columns ] attrs = df.filter(pl.col(DEFAULT_ATTR_KEYS.EDGE_ID) == parent_edge_id).drop(drop_cols).rows(named=True)[0] + # A rustworkx-family root shares its payload, EDGE_ID included; other backends + # hand back a plain dict, so stamp the root edge id explicitly (as `bulk_add_edges` + # does). Otherwise the local edge keeps the -1 placeholder and disagrees with + # `_edge_map_to_root`, so reads by edge id find no row. + attrs[DEFAULT_ATTR_KEYS.EDGE_ID] = parent_edge_id local_edge_id = self.rx_graph.add_edge( self._map_to_local(source_id), diff --git a/src/tracksdata/graph/_test/test_subgraph.py b/src/tracksdata/graph/_test/test_subgraph.py index f6993677..520bbeba 100644 --- a/src/tracksdata/graph/_test/test_subgraph.py +++ b/src/tracksdata/graph/_test/test_subgraph.py @@ -1648,6 +1648,23 @@ def test_add_edge_to_view_basic(graph_backend: BaseGraph) -> None: assert row["weight"].item() == 0.5 +def test_add_edge_to_view_keeps_edge_id(graph_backend: BaseGraph) -> None: + """A revived edge's local row must carry the root edge id, so reads by id work.""" + graph_backend.add_edge_attr_key("weight", pl.Float64) + + n0 = graph_backend.add_node({"t": 0}) + n1 = graph_backend.add_node({"t": 1}) + root_edge_id = graph_backend.add_edge(n0, n1, {"weight": 1.5}) + + view = graph_backend.filter().subgraph() + view.remove_edge_from_view(n0, n1) + view.add_edge_to_view(n0, n1) + + assert view.edge_id(n0, n1) == root_edge_id + assert view.edge_attrs(attr_keys=["weight"])[DEFAULT_ATTR_KEYS.EDGE_ID].to_list() == [root_edge_id] + assert view.edges[root_edge_id]["weight"] == 1.5 + + def test_add_edge_to_view_validation(graph_backend: BaseGraph) -> None: """Bad inputs raise ValueError; sync=False raises RuntimeError.""" graph_backend.add_node_attr_key("x", pl.Float64) From bdacacd7f809f9a388551133d6ee9f872832c3e8 Mon Sep 17 00:00:00 2001 From: Teun Huijben <45037215+TeunHuijben@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:00:25 -0700 Subject: [PATCH 6/6] fix np scalars becoming blobs in SQLGraph (#327) * fix np scalars becoming blobs in SQLGraph * remove unnecessary addition --- src/tracksdata/graph/_sql_graph.py | 36 ++++--- .../graph/_test/test_graph_backends.py | 95 +++++++++++++++++++ 2 files changed, 120 insertions(+), 11 deletions(-) diff --git a/src/tracksdata/graph/_sql_graph.py b/src/tracksdata/graph/_sql_graph.py index 151fb3e7..fe85dc8d 100644 --- a/src/tracksdata/graph/_sql_graph.py +++ b/src/tracksdata/graph/_sql_graph.py @@ -59,13 +59,21 @@ 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. + Parameters ---------- data : dict[str, Any] The data to convert. Modified in place. """ for k, v in data.items(): - if np.isscalar(v) and hasattr(v, "item"): + # `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() @@ -960,6 +968,11 @@ def _flatten_attrs_for_write( write path (``bulk_add_nodes``, ``bulk_add_edges``, ``update_node_attrs``, ``update_edge_attrs``) since the single-node/edge wrappers in :class:`BaseGraph` now delegate to the bulk variants. + + Numpy scalars are converted to their native Python counterparts here -- + including the ones nested inside a struct value -- so that every write path + hands the driver values it can map onto the column's declared dtype, + see :func:`_data_numpy_to_native`. """ result: dict[str, Any] = {} for key, value in attrs.items(): @@ -968,6 +981,7 @@ def _flatten_attrs_for_write( result.update(flatten_struct_value(key, value, schema.dtype)) else: result[key] = value + _data_numpy_to_native(result) return result def bulk_add_nodes( @@ -1018,7 +1032,12 @@ def bulk_add_nodes( node_ids = [] insert_rows = [] for i, node in enumerate(nodes): - time = node["t"] + # numpy times must be converted before the id arithmetic below, otherwise + # `time * node_id_time_multiplier` silently overflows for narrow dtypes + # (e.g. np.int32) and the resulting id is a numpy scalar itself. + time = node[DEFAULT_ATTR_KEYS.T] + if isinstance(time, np.generic): + time = time.item() if indices is None: default_node_id = (time * self.node_id_time_multiplier) - 1 @@ -1026,7 +1045,7 @@ def bulk_add_nodes( # Update max_id tracking only for auto-generated IDs self._max_id_per_time[time] = node_id else: - node_id = indices[i] + node_id = int(indices[i]) node_ids.append(node_id) insert_rows.append({**node, DEFAULT_ATTR_KEYS.NODE_ID: node_id}) @@ -1151,9 +1170,6 @@ def bulk_add_edges( return None edge_schemas = self._edge_attr_schemas() - for edge in edges: - _data_numpy_to_native(edge) - edges = [self._flatten_attrs_for_write(edge, edge_schemas) for edge in edges] if return_ids: @@ -1188,8 +1204,8 @@ def add_overlap( The ID of the added overlap. """ overlap = self.Overlap( - source_id=source_id, - target_id=target_id, + source_id=int(source_id), + target_id=int(target_id), ) with Session(self._engine) as session: session.add(overlap) @@ -1216,7 +1232,7 @@ def bulk_add_overlaps( if hasattr(overlaps, "tolist"): overlaps = overlaps.tolist() - overlaps = [{"source_id": source_id, "target_id": target_id} for source_id, target_id in overlaps] + overlaps = [{"source_id": int(source_id), "target_id": int(target_id)} for source_id, target_id in overlaps] self._chunked_sa_write(Session.bulk_insert_mappings, overlaps, self.Overlap) def overlaps( @@ -2004,8 +2020,6 @@ def _update_table( ids = ids.tolist() # Handle array values with bulk_update_mappings - attrs = attrs.copy() - _data_numpy_to_native(attrs) schemas = self._attr_schemas_for_table(table_class) attrs = self._flatten_attrs_for_write(attrs, schemas) diff --git a/src/tracksdata/graph/_test/test_graph_backends.py b/src/tracksdata/graph/_test/test_graph_backends.py index 15cb0a99..4712340b 100644 --- a/src/tracksdata/graph/_test/test_graph_backends.py +++ b/src/tracksdata/graph/_test/test_graph_backends.py @@ -133,6 +133,101 @@ def test_array_attr_read_honors_declared_dtype(graph_backend: BaseGraph) -> None assert nodes_df["values"].to_list() == [[50.0, 50.0], [1.5, 1.5]] +def test_add_node_and_edge_with_numpy_scalars(graph_backend: BaseGraph) -> None: + """Numpy scalars must be stored with the column's declared dtype, not as raw byte buffers. + + Indexing any value out of a numpy array yields a numpy scalar, so this is the + norm for importers (geff/CSV/...) feeding values into the graph. + """ + graph_backend.add_node_attr_key("val", dtype=pl.Int64, default_value=-1) + graph_backend.add_node_attr_key("pos", dtype=pl.Float64, default_value=0.0) + graph_backend.add_node_attr_key("flag", dtype=pl.Boolean, default_value=False) + graph_backend.add_edge_attr_key("weight", dtype=pl.Int64, default_value=0) + + node_1 = graph_backend.add_node( + {"t": np.int64(0), "val": np.int64(7), "pos": np.float32(1.5), "flag": np.bool_(True)} + ) + node_2, node_3 = graph_backend.bulk_add_nodes( + [ + {"t": np.int32(1), "val": np.int32(8), "pos": np.float64(2.5), "flag": np.bool_(False)}, + {"t": 2, "val": 9, "pos": 3.5, "flag": True}, + ] + ) + + nodes_df = graph_backend.node_attrs(attr_keys=["t", "val", "pos", "flag"]).sort("t") + assert nodes_df.schema["val"] == pl.Int64 + assert nodes_df["t"].to_list() == [0, 1, 2] + assert nodes_df["val"].to_list() == [7, 8, 9] + assert nodes_df["pos"].to_list() == [1.5, 2.5, 3.5] + assert nodes_df["flag"].to_list() == [True, False, True] + + graph_backend.add_edge(np.int64(node_1), np.int64(node_2), {"weight": np.int64(3)}) + graph_backend.bulk_add_edges( + [ + { + DEFAULT_ATTR_KEYS.EDGE_SOURCE: np.int64(node_2), + DEFAULT_ATTR_KEYS.EDGE_TARGET: np.int64(node_3), + "weight": np.int32(4), + } + ] + ) + + edges_df = graph_backend.edge_attrs( + attr_keys=[DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET, "weight"] + ).sort("weight") + assert edges_df["weight"].to_list() == [3, 4] + assert edges_df[DEFAULT_ATTR_KEYS.EDGE_SOURCE].to_list() == [node_1, node_2] + assert edges_df[DEFAULT_ATTR_KEYS.EDGE_TARGET].to_list() == [node_2, node_3] + + graph_backend.add_overlap(np.int64(node_1), np.int64(node_2)) + graph_backend.bulk_add_overlaps([[np.int64(node_2), np.int64(node_3)]]) + assert sorted(graph_backend.overlaps()) == sorted([[node_1, node_2], [node_2, node_3]]) + + +def test_sql_node_ids_from_narrow_numpy_time() -> None: + """A narrow numpy `t` must not overflow the `t * node_id_time_multiplier` id arithmetic.""" + graph = SQLGraph( + drivername="sqlite", + database=":memory:", + engine_kwargs={"connect_args": {"check_same_thread": False}}, + ) + # np.int32(3) * 1_000_000_000 wraps around to a negative number in int32 arithmetic + node_ids = graph.bulk_add_nodes([{"t": np.int32(3)}, {"t": np.int32(3)}]) + + assert node_ids == [3 * graph.node_id_time_multiplier, 3 * graph.node_id_time_multiplier + 1] + assert all(isinstance(node_id, int) for node_id in node_ids) + assert graph.node_ids() == node_ids + + +def test_add_node_with_numpy_scalars_in_struct(graph_backend: BaseGraph) -> None: + """Numpy scalars nested inside a struct attribute must also honor the declared dtype.""" + graph_backend.add_node_attr_key("m", dtype=pl.Struct({"a": pl.Int64, "b": pl.Float64})) + + graph_backend.add_node({"t": 0, "m": {"a": np.int64(3), "b": np.float64(0.25)}}) + graph_backend.bulk_add_nodes([{"t": 1, "m": {"a": np.int32(4), "b": np.float32(0.5)}}]) + + nodes_df = graph_backend.node_attrs(attr_keys=["t", "m"]).sort("t") + assert nodes_df["m"].to_list() == [{"a": 3, "b": 0.25}, {"a": 4, "b": 0.5}] + + +def test_update_attrs_with_numpy_scalars(graph_backend: BaseGraph) -> None: + """The update path must coerce numpy scalars just like the insert path.""" + graph_backend.add_node_attr_key("val", dtype=pl.Int64, default_value=-1) + graph_backend.add_edge_attr_key("weight", dtype=pl.Int64, default_value=0) + + node_1 = graph_backend.add_node({"t": 0, "val": 0}) + node_2 = graph_backend.add_node({"t": 1, "val": 0}) + edge_id = graph_backend.add_edge(node_1, node_2, {"weight": 0}) + + graph_backend.update_node_attrs(attrs={"val": np.int64(5)}, node_ids=[node_1]) + graph_backend.update_node_attrs(attrs={"val": [np.int32(6)]}, node_ids=[node_2]) + graph_backend.update_edge_attrs(attrs={"weight": np.int64(7)}, edge_ids=[edge_id]) + + nodes_df = graph_backend.node_attrs(attr_keys=["t", "val"]).sort("t") + assert nodes_df["val"].to_list() == [5, 6] + assert graph_backend.edge_attrs(attr_keys=["weight"])["weight"].to_list() == [7] + + def test_remove_edge_by_id(graph_backend: BaseGraph) -> None: """Test removing an edge by ID across backends using unified API.""" # Setup