Refactoring node and edge attribute key API by JoOkuma · Pull Request #244 · royerlab/tracksdata

JoOkuma · 2026-01-20T18:27:29Z

Closes #193 and #243

- Add AttrSchema dataclass with defensive copying - Create process_attr_key_args() helper to eliminate duplication - Refactor dtype inference and SQLAlchemy mapping to use dictionaries - Add overloaded method signatures for convenience and schema modes - Replace _node/edge_attr_keys lists with _attr_schemas dicts - Add comprehensive test suite (16 tests passing) Breaking change: dtype parameter now required. Next phase will update all calling sites throughout the codebase.

…add_edge_attr_key This is a breaking change that requires dtype (polars.DataType) to be explicitly specified when adding attribute keys to graphs. Core changes: - Add AttrSchema dataclass to store key, dtype, and default_value together - Add utility functions: infer_default_value_from_dtype, polars_dtype_to_sqlalchemy_type, validate_default_value_dtype_compatibility, and process_attr_key_args - Update BaseGraph abstract methods with overloaded signatures supporting both direct parameters and AttrSchema objects - Replace _node_attr_keys/_edge_attr_keys lists with _node_attr_schemas dicts in RustWorkXGraph and SQLGraph - Remove _boolean_columns tracking in SQLGraph (now uses AttrSchema) - Fix _polars_schema_override to use == instead of isinstance for pl.Boolean check Backend implementations: - RustWorkXGraph: Use process_attr_key_args helper, store AttrSchema objects - SQLGraph: Use process_attr_key_args helper, pass AttrSchema to _add_new_column - GraphView: Delegate to root graph with all parameters preserved - Add dtype inference fallback for complex objects that polars can't parse Updated all callers: - Graph methods: from_other(), match() - Node operators: RandomNodes, RegionPropsNodes, GenericNodesOperator - Edge operators: DistanceEdges, GenericEdgesOperator - Solvers: ILPSolver, NearestNeighborsSolver - IO: numpy_array, ctc

Update all test files to use the new required dtype parameter for add_node_attr_key and add_edge_attr_key methods. Key test fixes: - Add dtype inference logic in test_add_node_attr_key based on value type - Fix bbox attribute calls to use pl.Array(pl.Int64, size) instead of passing numpy arrays as dtype - Add missing polars imports where needed - Fix argument order from (key, False, dtype=pl.Boolean) to (key, pl.Boolean, default_value=False)

codecov-commenter · 2026-01-20T18:31:43Z

Codecov Report

❌ Patch coverage is 91.00529% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.15%. Comparing base (574694d) to head (1cb217f).

Files with missing lines	Patch %	Lines
src/tracksdata/graph/_sql_graph.py	85.00%	4 Missing and 2 partials ⚠️
src/tracksdata/graph/_graph_view.py	73.33%	2 Missing and 2 partials ⚠️
src/tracksdata/graph/_rustworkx_graph.py	90.24%	3 Missing and 1 partial ⚠️
src/tracksdata/utils/_dtypes.py	96.61%	1 Missing and 1 partial ⚠️
src/tracksdata/nodes/_mask.py	80.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #244      +/-   ##
==========================================
- Coverage   88.35%   88.15%   -0.21%     
==========================================
  Files          56       56              
  Lines        4104     4203      +99     
  Branches      714      733      +19     
==========================================
+ Hits         3626     3705      +79     
- Misses        291      311      +20     
  Partials      187      187

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

JoOkuma · 2026-01-20T19:31:59Z

Hi @TeunHuijben and @yfukai,

I would appreciate it if you both could review this PR, at least the public API, since it involves breaking changes.
Before you do that, I need to give it a pass because a big part of it was generated by an LLM.
I'll tag you again once I'm done.

JoOkuma · 2026-01-20T20:55:57Z

src/tracksdata/graph/_rustworkx_graph.py

        rx_graph = self.rx_graph
        for node_id in rx_graph.node_indices():
-            rx_graph[node_id][key] = default_value
+            rx_graph[node_id][schema.key] = schema.default_value


Do we really need to add the default_value?
This could be solved when returning the attrs value.

Yeah I think returning the attrs value for unset attribute would be better.

JoOkuma · 2026-01-20T20:56:10Z

src/tracksdata/graph/_rustworkx_graph.py

+        # Add to all existing edges
        for edge_attr in self.rx_graph.edges():
-            edge_attr[key] = default_value
+            edge_attr[schema.key] = schema.default_value


Same as above.

TeunHuijben · 2026-01-21T00:21:58Z

Haven't personally tested the code, but the API looks good!

JoOkuma · 2026-01-21T00:25:06Z

Hi @yfukai , this should be ready for review.
I'm going to open a separate PR that stops setting the attributes of existing nodes when adding a new key, and it will happen when querying node|edge attributes.

TeunHuijben · 2026-01-23T19:03:14Z

Hi @JoOkuma and @yfukai, what is the status of this PR, I would really like the default-feature PR to be merged, which depends on this one 😇

yfukai · 2026-01-23T22:56:40Z

Hi sorry @TeunHuijben I'm still reviewing it, I'll finish that within this weekend!

TeunHuijben · 2026-01-23T23:06:55Z

Hi @yfukai, no worries, take your time! Was just wondering whether your thumbs-up meant approval or whether you were reviewing. Thank you! 😊

yfukai

This looks much nicer than what we currently have, thanks @JoOkuma!
I think this PR is good to merge, but

It would be even better if the dtype does not change with a round-trip with from_other. For this purpose, maybe we can store the AttrSchemas into a read-only metadata.
~~For RustworkXGraph, maybe we can skip setting the default value to the rx_graph and return the default value for unset attr to save the initialization time.~~ → I found this addressed in the next PR, thanks!

yfukai · 2026-01-24T07:36:21Z

src/tracksdata/graph/_rustworkx_graph.py

        rx_graph = self.rx_graph
        for node_id in rx_graph.node_indices():
-            rx_graph[node_id][key] = default_value
+            rx_graph[node_id][schema.key] = schema.default_value


Yeah I think returning the attrs value for unset attribute would be better.

JoOkuma · 2026-01-29T00:52:10Z

Awesome! Thanks @TeunHuijben and @yfukai

@yfukai, I added checking if the schema matches when calling from_other
I noticed that the number of edges and overlaps is limited to 32 bits due to SQLite's autoincrement, just so you keep this in mind in case you run into a weird bug in the future.

Can you also take a quick look at #246 if you've time?
No is a valid answer if you're busy. I'm sure you have other things going on.

JoOkuma added 3 commits January 20, 2026 09:47

JoOkuma added 4 commits January 20, 2026 10:54

fixing correct usage of dtypes and using default when possible

5392cbf

fixing more attr usage

b2368e1

removing unused function

e3ac6ea

improving attr key default usage

b7b00f4

improving testing

2bcafcb

JoOkuma commented Jan 20, 2026

View reviewed changes

JoOkuma added 5 commits January 20, 2026 13:03

fixing spatial filtering typing

8288a69

simplifying sql type detection case

58256e2

refactoring generic usage

a05788c

moving tests and fixing get attr usage

b83b587

cleanup on sql type conversion

3ca9164

fixing docs

7940fee

JoOkuma mentioned this pull request Jan 22, 2026

Delayed default value setting #246

Open

Merge branch 'main' into jookuma/attr-key-refactoring

1cb217f

yfukai reviewed Jan 24, 2026

View reviewed changes

JoOkuma merged commit c5d40b2 into royerlab:main Jan 29, 2026
7 checks passed

JoOkuma deleted the jookuma/attr-key-refactoring branch January 29, 2026 00:49

Conversation

JoOkuma commented Jan 20, 2026

Uh oh!

codecov-commenter commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

JoOkuma commented Jan 20, 2026

Uh oh!

JoOkuma Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

yfukai Jan 24, 2026

Choose a reason for hiding this comment

Uh oh!

JoOkuma Jan 20, 2026

Choose a reason for hiding this comment

Uh oh!

TeunHuijben commented Jan 21, 2026

Uh oh!

JoOkuma commented Jan 21, 2026

Uh oh!

TeunHuijben commented Jan 23, 2026

Uh oh!

yfukai commented Jan 23, 2026

Uh oh!

TeunHuijben commented Jan 23, 2026

Uh oh!

yfukai left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yfukai Jan 24, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

JoOkuma commented Jan 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

codecov-commenter commented Jan 20, 2026 •

edited

Loading

yfukai left a comment •

edited

Loading