Changed successors and predecessors so that they return IDs.#188
Changed successors and predecessors so that they return IDs.#188JoOkuma merged 7 commits intoroyerlab:mainfrom
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #188 +/- ##
==========================================
- Coverage 88.17% 88.05% -0.13%
==========================================
Files 51 51
Lines 3612 3632 +20
Branches 625 633 +8
==========================================
+ Hits 3185 3198 +13
- Misses 256 261 +5
- Partials 171 173 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
This is also ready! @JoOkuma |
…cessor_predecessor_nodeids
| if isinstance(attr_keys, str): | ||
| attr_keys = [attr_keys] | ||
| valid_schema = None | ||
| neighbors: dict[int, pl.DataFrame] = {} | ||
| for node_id in node_ids: | ||
| neighbors_indices = neighbors_func(rx_graph, node_id) | ||
| neighbors_data: list[dict[str, Any]] = [rx_graph[i] for i in neighbors_indices] | ||
|
|
||
| if attr_keys is not None: | ||
| neighbors_data = [ | ||
| {k: edge_data[k] for k in attr_keys if k != DEFAULT_ATTR_KEYS.NODE_ID} | ||
| for edge_data in neighbors_data | ||
| ] | ||
|
|
||
| if attr_keys is not None: | ||
| neighbors_data = [ | ||
| {k: edge_data[k] for k in attr_keys if k != DEFAULT_ATTR_KEYS.NODE_ID} | ||
| for edge_data in neighbors_data | ||
| ] | ||
|
|
||
| if len(neighbors_data) > 0: | ||
| df = pl.DataFrame(neighbors_data) | ||
| if attr_keys is None or DEFAULT_ATTR_KEYS.NODE_ID in attr_keys: | ||
| df = df.with_columns( | ||
| pl.Series(DEFAULT_ATTR_KEYS.NODE_ID, np.asarray(neighbors_indices, dtype=int)), | ||
| ) | ||
| neighbors[node_id] = df | ||
| valid_schema = neighbors[node_id].schema | ||
| if len(neighbors_data) > 0: | ||
| df = pl.DataFrame(neighbors_data) | ||
| if attr_keys is None or DEFAULT_ATTR_KEYS.NODE_ID in attr_keys: | ||
| df = df.with_columns( | ||
| pl.Series(DEFAULT_ATTR_KEYS.NODE_ID, np.asarray(neighbors_indices, dtype=int)), | ||
| ) | ||
| neighbors[node_id] = df | ||
| valid_schema = neighbors[node_id].schema |
There was a problem hiding this comment.
I think we could reduce the amount of code duplication when return_attrs=True by querying the indices with a call of the same function with return_attrs=False and then querying the dataframes. What do you think?
There was a problem hiding this comment.
Yeah that was also an option which is better maybe. Let me try on that track!
There was a problem hiding this comment.
LGTM, I just renamed a variable. Once the test passes, I'll merge, thanks for the PR!
Co-authored-by: Jordão Bragantini <jordao.bragantini@gmail.com>
|
Hi @JoOkuma, this may not be exactly the same as what you meant but I tried to reduce the code duplicates! Could you review the udpated code? |
Closes #187.
API and Interface Updates
return_attrsparameter (defaulting toFalse) tosuccessorsandpredecessorsmethods insrc/tracksdata/graph/_base_graph.py, updating their overloads and documentation to clarify return types for both node IDs and attributes._rustworkx_graph.py,_sql_graph.py,_graph_view.py) to propagate thereturn_attrsparameter and handle both return modes, including changes to_get_neighborsand related methods.Internal Logic and Performance
return_attrsisFalse, only node IDs are returned, avoiding unnecessary DataFrame creation and improving efficiency for large graphs.