Skip to content

Commit 79c64ab

Browse files
FlorentinDFlorentinD
committed
Add steinerTree to v2 endpoints
1 parent 9301b15 commit 79c64ab

File tree

9 files changed

+1029
-3
lines changed

9 files changed

+1029
-3
lines changed

‎graphdatascience/procedure_surface/api/pathfinding/steiner_tree_endpoints.py‎

Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,251 @@
1+
from __future__ import annotations
2+
3+
from abc import ABC, abstractmethod
4+
from typing import Any
5+
6+
from pandas import DataFrame
7+
8+
from graphdatascience.procedure_surface.api.base_result import BaseResult
9+
from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
10+
from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
11+
12+
13+
class SteinerTreeMutateResult(BaseResult):
14+
relationships_written: int
15+
mutate_millis: int
16+
effective_node_count: int
17+
effective_target_nodes_count: int
18+
total_weight: float
19+
pre_processing_millis: int
20+
compute_millis: int
21+
configuration: dict[str, Any]
22+
23+
24+
class SteinerTreeWriteResult(BaseResult):
25+
relationships_written: int
26+
write_millis: int
27+
effective_node_count: int
28+
effective_target_nodes_count: int
29+
total_weight: float
30+
pre_processing_millis: int
31+
compute_millis: int
32+
configuration: dict[str, Any]
33+
34+
35+
class SteinerTreeStatsResult(BaseResult):
36+
effective_node_count: int
37+
effective_target_nodes_count: int
38+
total_weight: float
39+
pre_processing_millis: int
40+
compute_millis: int
41+
configuration: dict[str, Any]
42+
43+
44+
class SteinerTreeEndpoints(ABC):
45+
@abstractmethod
46+
def stream(
47+
self,
48+
G: GraphV2,
49+
source_node: int,
50+
target_nodes: list[int],
51+
relationship_weight_property: str | None = None,
52+
delta: float = 2.0,
53+
apply_rerouting: bool = False,
54+
relationship_types: list[str] | None = None,
55+
node_labels: list[str] | None = None,
56+
sudo: bool = False,
57+
log_progress: bool = True,
58+
username: str | None = None,
59+
concurrency: int | None = None,
60+
job_id: str | None = None,
61+
) -> DataFrame:
62+
"""
63+
Runs the Steiner tree algorithm and returns the result as a DataFrame.
64+
65+
Parameters
66+
----------
67+
G : GraphV2
68+
The graph to run the algorithm on.
69+
source_node : int
70+
The source node (root) for the Steiner tree.
71+
target_nodes : list[int]
72+
The list of target nodes (terminals) that must be connected.
73+
relationship_weight_property : str, optional
74+
The name of the relationship property to use as weights.
75+
delta : float, default=2.0
76+
The delta parameter for the shortest path computation used internally.
77+
apply_rerouting : bool, default=False
78+
Whether to apply rerouting optimization to improve the tree.
79+
relationship_types : list[str], optional
80+
Filter to only use relationships of specific types.
81+
node_labels : list[str], optional
82+
Filter to only use nodes with specific labels.
83+
sudo : bool, default=False
84+
Whether to run with elevated privileges.
85+
log_progress : bool, default=True
86+
Whether to log progress during execution.
87+
username : str, optional
88+
The username to use for logging.
89+
concurrency : int, optional
90+
The number of threads to use for parallel computation.
91+
job_id : str, optional
92+
An optional job ID for tracking the operation.
93+
94+
Returns
95+
-------
96+
DataFrame
97+
A DataFrame containing the edges in the computed Steiner tree.
98+
"""
99+
...
100+
101+
@abstractmethod
102+
def stats(
103+
self,
104+
G: GraphV2,
105+
source_node: int,
106+
target_nodes: list[int],
107+
relationship_weight_property: str | None = None,
108+
delta: float = 2.0,
109+
apply_rerouting: bool = False,
110+
relationship_types: list[str] | None = None,
111+
node_labels: list[str] | None = None,
112+
sudo: bool = False,
113+
log_progress: bool = True,
114+
username: str | None = None,
115+
concurrency: int | None = None,
116+
job_id: str | None = None,
117+
) -> SteinerTreeStatsResult:
118+
"""
119+
Runs the Steiner tree algorithm in stats mode, returning statistics without modifying the graph.
120+
121+
Returns
122+
-------
123+
SteinerTreeStatsResult
124+
Statistics about the computed Steiner tree.
125+
"""
126+
...
127+
128+
@abstractmethod
129+
def mutate(
130+
self,
131+
G: GraphV2,
132+
mutate_relationship_type: str,
133+
mutate_property: str,
134+
source_node: int,
135+
target_nodes: list[int],
136+
relationship_weight_property: str | None = None,
137+
delta: float = 2.0,
138+
apply_rerouting: bool = False,
139+
relationship_types: list[str] | None = None,
140+
node_labels: list[str] | None = None,
141+
sudo: bool = False,
142+
log_progress: bool = True,
143+
username: str | None = None,
144+
concurrency: int | None = None,
145+
job_id: str | None = None,
146+
) -> SteinerTreeMutateResult:
147+
"""
148+
Runs the Steiner tree algorithm and adds the result as new relationships to the in-memory graph.
149+
150+
Parameters
151+
----------
152+
mutate_relationship_type : str
153+
The relationship type to use for the new relationships.
154+
mutate_property : str
155+
The property name to store the edge weight.
156+
157+
Returns
158+
-------
159+
SteinerTreeMutateResult
160+
Result containing statistics and timing information.
161+
"""
162+
...
163+
164+
@abstractmethod
165+
def write(
166+
self,
167+
G: GraphV2,
168+
write_relationship_type: str,
169+
write_property: str,
170+
source_node: int,
171+
target_nodes: list[int],
172+
relationship_weight_property: str | None = None,
173+
delta: float = 2.0,
174+
apply_rerouting: bool = False,
175+
relationship_types: list[str] | None = None,
176+
node_labels: list[str] | None = None,
177+
sudo: bool = False,
178+
log_progress: bool = True,
179+
username: str | None = None,
180+
concurrency: int | None = None,
181+
job_id: str | None = None,
182+
write_concurrency: int | None = None,
183+
) -> SteinerTreeWriteResult:
184+
"""
185+
Runs the Steiner tree algorithm and writes the result back to the Neo4j database.
186+
187+
Parameters
188+
----------
189+
write_relationship_type : str
190+
The relationship type to use for the new relationships.
191+
write_property : str
192+
The property name to store the edge weight.
193+
write_concurrency : int, optional
194+
The number of threads to use for writing results.
195+
196+
Returns
197+
-------
198+
SteinerTreeWriteResult
199+
Result containing statistics and timing information.
200+
"""
201+
...
202+
203+
@abstractmethod
204+
def estimate(
205+
self,
206+
G: GraphV2 | dict[str, Any],
207+
source_node: int,
208+
target_nodes: list[int],
209+
relationship_weight_property: str | None = None,
210+
delta: float = 2.0,
211+
apply_rerouting: bool = False,
212+
relationship_types: list[str] | None = None,
213+
node_labels: list[str] | None = None,
214+
sudo: bool = False,
215+
username: str | None = None,
216+
concurrency: int | None = None,
217+
) -> EstimationResult:
218+
"""
219+
Estimates the memory requirements for running the Steiner tree algorithm.
220+
221+
Parameters
222+
----------
223+
G : GraphV2 | dict[str, Any]
224+
The graph to estimate for, or a dictionary with nodeCount and relationshipCount.
225+
source_node : int
226+
The source node (root) for the Steiner tree.
227+
target_nodes : list[int]
228+
The list of target nodes (terminals) that must be connected.
229+
relationship_weight_property : str, optional
230+
The name of the relationship property to use as weights.
231+
delta : float, default=2.0
232+
The delta parameter for the shortest path computation.
233+
apply_rerouting : bool, default=False
234+
Whether to apply rerouting optimization.
235+
relationship_types : list[str], optional
236+
Filter to only use relationships of specific types.
237+
node_labels : list[str], optional
238+
Filter to only use nodes with specific labels.
239+
sudo : bool, default=False
240+
Whether to run with elevated privileges.
241+
username : str, optional
242+
The username to use for logging.
243+
concurrency : int, optional
244+
The number of threads to use for parallel computation.
245+
246+
Returns
247+
-------
248+
EstimationResult
249+
Memory estimation results including required bytes and percentages.
250+
"""
251+
...

‎graphdatascience/procedure_surface/arrow/pathfinding/single_source_bellman_ford_arrow_endpoints.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
from __future__ import annotations
22

33
from typing import Any
4-
from unittest import result
54

65
from pandas import DataFrame
76

0 commit comments

Comments
 (0)