Skip to content

Commit 9274e63

Browse files
committed
Map spanning tree to v2 endpoints
1 parent 47c9332 commit 9274e63

File tree

7 files changed

+974
-4
lines changed

7 files changed

+974
-4
lines changed
Lines changed: 299 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,299 @@
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 SpanningTreeMutateResult(BaseResult):
14+
relationships_written: int
15+
mutate_millis: int
16+
effective_node_count: int
17+
total_weight: float
18+
pre_processing_millis: int
19+
compute_millis: int
20+
configuration: dict[str, Any]
21+
22+
23+
class SpanningTreeWriteResult(BaseResult):
24+
relationships_written: int
25+
write_millis: int
26+
effective_node_count: int
27+
total_weight: float
28+
pre_processing_millis: int
29+
compute_millis: int
30+
configuration: dict[str, Any]
31+
32+
33+
class SpanningTreeStatsResult(BaseResult):
34+
effective_node_count: int
35+
total_weight: float
36+
pre_processing_millis: int
37+
compute_millis: int
38+
configuration: dict[str, Any]
39+
40+
41+
class SpanningTreeEndpoints(ABC):
42+
@abstractmethod
43+
def stream(
44+
self,
45+
G: GraphV2,
46+
source_node: int,
47+
relationship_weight_property: str | None = None,
48+
objective: str = "minimum",
49+
relationship_types: list[str] | None = None,
50+
node_labels: list[str] | None = None,
51+
sudo: bool = False,
52+
log_progress: bool = True,
53+
username: str | None = None,
54+
concurrency: int | None = None,
55+
job_id: str | None = None,
56+
) -> DataFrame:
57+
"""
58+
Runs the Spanning tree algorithm and returns the result as a DataFrame.
59+
60+
Parameters
61+
----------
62+
G : GraphV2
63+
The graph to run the algorithm on.
64+
source_node : int
65+
The source node (root) for the Spanning tree.
66+
relationship_weight_property : str, optional
67+
The name of the relationship property to use as weights.
68+
objective : str, default="minimum"
69+
The objective function to optimize. Either "minimum" or "maximum".
70+
relationship_types : list[str], optional
71+
Filter to only use relationships of specific types.
72+
node_labels : list[str], optional
73+
Filter to only use nodes with specific labels.
74+
sudo : bool, default=False
75+
Whether to run with elevated privileges.
76+
log_progress : bool, default=True
77+
Whether to log progress during execution.
78+
username : str, optional
79+
The username to use for logging.
80+
concurrency : int, optional
81+
The number of threads to use for parallel computation.
82+
job_id : str, optional
83+
An optional job ID for tracking the operation.
84+
85+
Returns
86+
-------
87+
DataFrame
88+
A DataFrame containing the edges in the computed Spanning tree.
89+
"""
90+
...
91+
92+
@abstractmethod
93+
def stats(
94+
self,
95+
G: GraphV2,
96+
source_node: int,
97+
relationship_weight_property: str | None = None,
98+
objective: str = "minimum",
99+
relationship_types: list[str] | None = None,
100+
node_labels: list[str] | None = None,
101+
sudo: bool = False,
102+
log_progress: bool = True,
103+
username: str | None = None,
104+
concurrency: int | None = None,
105+
job_id: str | None = None,
106+
) -> SpanningTreeStatsResult:
107+
"""
108+
Runs the Spanning tree algorithm in stats mode, returning statistics without modifying the graph.
109+
110+
Parameters
111+
----------
112+
G : GraphV2
113+
The graph to run the algorithm on.
114+
source_node : int
115+
The source node (root) for the Spanning tree.
116+
relationship_weight_property : str, optional
117+
The name of the relationship property to use as weights.
118+
objective : str, default="minimum"
119+
The objective function to optimize. Either "minimum" or "maximum".
120+
relationship_types : list[str], optional
121+
Filter to only use relationships of specific types.
122+
node_labels : list[str], optional
123+
Filter to only use nodes with specific labels.
124+
sudo : bool, default=False
125+
Whether to run with elevated privileges.
126+
log_progress : bool, default=True
127+
Whether to log progress during execution.
128+
username : str, optional
129+
The username to use for logging.
130+
concurrency : int, optional
131+
The number of threads to use for parallel computation.
132+
job_id : str, optional
133+
An optional job ID for tracking the operation.
134+
135+
Returns
136+
-------
137+
SpanningTreeStatsResult
138+
Statistics about the computed Spanning tree.
139+
"""
140+
...
141+
142+
@abstractmethod
143+
def mutate(
144+
self,
145+
G: GraphV2,
146+
mutate_relationship_type: str,
147+
mutate_property: str,
148+
source_node: int,
149+
relationship_weight_property: str | None = None,
150+
objective: str = "minimum",
151+
relationship_types: list[str] | None = None,
152+
node_labels: list[str] | None = None,
153+
sudo: bool = False,
154+
log_progress: bool = True,
155+
username: str | None = None,
156+
concurrency: int | None = None,
157+
job_id: str | None = None,
158+
) -> SpanningTreeMutateResult:
159+
"""
160+
Runs the Spanning tree algorithm and adds the result as new relationships to the in-memory graph.
161+
162+
Parameters
163+
----------
164+
G : GraphV2
165+
The graph to run the algorithm on.
166+
mutate_relationship_type : str
167+
The relationship type to use for the new relationships.
168+
mutate_property : str
169+
The property name to store the edge weight.
170+
source_node : int
171+
The source node (root) for the Spanning tree.
172+
relationship_weight_property : str, optional
173+
The name of the relationship property to use as weights.
174+
objective : str, default="minimum"
175+
The objective function to optimize. Either "minimum" or "maximum".
176+
relationship_types : list[str], optional
177+
Filter to only use relationships of specific types.
178+
node_labels : list[str], optional
179+
Filter to only use nodes with specific labels.
180+
sudo : bool, default=False
181+
Whether to run with elevated privileges.
182+
log_progress : bool, default=True
183+
Whether to log progress during execution.
184+
username : str, optional
185+
The username to use for logging.
186+
concurrency : int, optional
187+
The number of threads to use for parallel computation.
188+
job_id : str, optional
189+
An optional job ID for tracking the operation.
190+
191+
Returns
192+
-------
193+
SpanningTreeMutateResult
194+
Result containing statistics and timing information.
195+
"""
196+
...
197+
198+
@abstractmethod
199+
def write(
200+
self,
201+
G: GraphV2,
202+
write_relationship_type: str,
203+
write_property: str,
204+
source_node: int,
205+
relationship_weight_property: str | None = None,
206+
objective: str = "minimum",
207+
relationship_types: list[str] | None = None,
208+
node_labels: list[str] | None = None,
209+
sudo: bool = False,
210+
log_progress: bool = True,
211+
username: str | None = None,
212+
concurrency: int | None = None,
213+
job_id: str | None = None,
214+
write_concurrency: int | None = None,
215+
) -> SpanningTreeWriteResult:
216+
"""
217+
Runs the Spanning tree algorithm and writes the result back to the Neo4j database.
218+
219+
Parameters
220+
----------
221+
G : GraphV2
222+
The graph to run the algorithm on.
223+
write_relationship_type : str
224+
The relationship type to use for the new relationships.
225+
write_property : str
226+
The property name to store the edge weight.
227+
source_node : int
228+
The source node (root) for the Spanning tree.
229+
relationship_weight_property : str, optional
230+
The name of the relationship property to use as weights.
231+
objective : str, default="minimum"
232+
The objective function to optimize. Either "minimum" or "maximum".
233+
relationship_types : list[str], optional
234+
Filter to only use relationships of specific types.
235+
node_labels : list[str], optional
236+
Filter to only use nodes with specific labels.
237+
sudo : bool, default=False
238+
Whether to run with elevated privileges.
239+
log_progress : bool, default=True
240+
Whether to log progress during execution.
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+
job_id : str, optional
246+
An optional job ID for tracking the operation.
247+
write_concurrency : int, optional
248+
The number of threads to use for writing results.
249+
250+
Returns
251+
-------
252+
SpanningTreeWriteResult
253+
Result containing statistics and timing information.
254+
"""
255+
...
256+
257+
@abstractmethod
258+
def estimate(
259+
self,
260+
G: GraphV2 | dict[str, Any],
261+
source_node: int,
262+
relationship_weight_property: str | None = None,
263+
objective: str = "minimum",
264+
relationship_types: list[str] | None = None,
265+
node_labels: list[str] | None = None,
266+
sudo: bool = False,
267+
username: str | None = None,
268+
concurrency: int | None = None,
269+
) -> EstimationResult:
270+
"""
271+
Estimates the memory requirements for running the Spanning tree algorithm.
272+
273+
Parameters
274+
----------
275+
G : GraphV2 | dict[str, Any]
276+
The graph to estimate for, or a dictionary with nodeCount and relationshipCount.
277+
source_node : int
278+
The source node (root) for the Spanning tree.
279+
relationship_weight_property : str, optional
280+
The name of the relationship property to use as weights.
281+
objective : str, default="minimum"
282+
The objective function to optimize. Either "minimum" or "maximum".
283+
relationship_types : list[str], optional
284+
Filter to only use relationships of specific types.
285+
node_labels : list[str], optional
286+
Filter to only use nodes with specific labels.
287+
sudo : bool, default=False
288+
Whether to run with elevated privileges.
289+
username : str, optional
290+
The username to use for logging.
291+
concurrency : int, optional
292+
The number of threads to use for parallel computation.
293+
294+
Returns
295+
-------
296+
EstimationResult
297+
Memory estimation results including required bytes and percentages.
298+
"""
299+
...

0 commit comments

Comments
 (0)