Skip to content

Commit 9b4829a

Browse files
DarthMaxDarthMax
authored
Merge pull request #925 from FlorentinD/v2endpoints-smaller-cleanups
Smaller cleanups for v2 endpoints
2 parents 30bcd60 + cd87a70 commit 9b4829a

19 files changed

+37
-51
lines changed

‎graphdatascience/arrow_client/v2/api_types.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,5 +12,6 @@ class JobStatus(ArrowBaseModel):
1212

1313

1414
class MutateResult(ArrowBaseModel):
15+
mutate_millis: int
1516
node_properties_written: int
1617
relationships_written: int

‎graphdatascience/arrow_client/v2/mutation_client.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import json
2+
import math
3+
import time
24

35
from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
46
from graphdatascience.arrow_client.v2.api_types import MutateResult
@@ -12,5 +14,7 @@ class MutationClient:
1214
def mutate_node_property(client: AuthenticatedArrowClient, job_id: str, mutate_property: str) -> MutateResult:
1315
mutate_config = {"jobId": job_id, "mutateProperty": mutate_property}
1416
encoded_config = json.dumps(mutate_config).encode("utf-8")
17+
start_time = time.time()
1518
mutate_arrow_res = client.do_action_with_retry(MutationClient.MUTATE_ENDPOINT, encoded_config)
16-
return MutateResult(**deserialize_single(mutate_arrow_res))
19+
mutate_millis = math.ceil((time.time() - start_time) * 1000)
20+
return MutateResult(mutateMillis=mutate_millis, **deserialize_single(mutate_arrow_res))

‎graphdatascience/procedure_surface/api/k1coloring_endpoints.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,6 @@ def write(
178178
concurrency: Optional[Any] = None,
179179
job_id: Optional[Any] = None,
180180
write_concurrency: Optional[Any] = None,
181-
write_to_result_store: Optional[bool] = None,
182181
min_community_size: Optional[int] = None,
183182
) -> K1ColoringWriteResult:
184183
"""

‎graphdatascience/procedure_surface/api/louvain_endpoints.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,7 +228,6 @@ def write(
228228
consecutive_ids: Optional[bool] = None,
229229
relationship_weight_property: Optional[str] = None,
230230
write_concurrency: Optional[Any] = None,
231-
write_to_result_store: Optional[bool] = None,
232231
min_community_size: Optional[int] = None,
233232
) -> LouvainWriteResult:
234233
"""

‎graphdatascience/procedure_surface/api/scc_endpoints.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,6 @@ def write(
165165
job_id: Optional[Any] = None,
166166
consecutive_ids: Optional[bool] = None,
167167
write_concurrency: Optional[Any] = None,
168-
write_to_result_store: Optional[bool] = None,
169168
) -> SccWriteResult:
170169
"""
171170
Executes the SCC algorithm and writes the results to the Neo4j database.

‎graphdatascience/procedure_surface/arrow/k1coloring_arrow_endpoints.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,12 +46,12 @@ def mutate(
4646
username=username,
4747
)
4848

49-
computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
49+
computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
5050
"v2/community.k1coloring", G, config, mutate_property
5151
)
5252

53-
computation_result["nodeCount"] = node_properties_written
54-
computation_result["mutateMillis"] = 0
53+
computation_result["nodeCount"] = mutate_result.node_properties_written
54+
computation_result["mutateMillis"] = mutate_result.mutate_millis
5555

5656
return K1ColoringMutateResult(**computation_result)
5757

@@ -129,7 +129,6 @@ def write(
129129
concurrency: Optional[int] = None,
130130
job_id: Optional[str] = None,
131131
write_concurrency: Optional[int] = None,
132-
write_to_result_store: Optional[bool] = None,
133132
min_community_size: Optional[int] = None,
134133
) -> K1ColoringWriteResult:
135134
config = self._node_property_endpoints.create_base_config(

‎graphdatascience/procedure_surface/arrow/kcore_arrow_endpoints.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,12 +36,12 @@ def mutate(
3636
sudo=sudo,
3737
)
3838

39-
computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
39+
computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
4040
"v2/community.kcore", G, config, mutate_property
4141
)
4242

43-
computation_result["nodePropertiesWritten"] = node_properties_written
44-
computation_result["mutateMillis"] = 0
43+
computation_result["nodePropertiesWritten"] = mutate_result.node_properties_written
44+
computation_result["mutateMillis"] = mutate_result.mutate_millis
4545

4646
return KCoreMutateResult(**computation_result)
4747

‎graphdatascience/procedure_surface/arrow/louvain_arrow_endpoints.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -50,12 +50,12 @@ def mutate(
5050
tolerance=tolerance,
5151
)
5252

53-
computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
53+
computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
5454
"v2/community.louvain", G, config, mutate_property
5555
)
5656

57-
computation_result["nodePropertiesWritten"] = node_properties_written
58-
computation_result["mutateMillis"] = 0
57+
computation_result["nodePropertiesWritten"] = mutate_result.node_properties_written
58+
computation_result["mutateMillis"] = mutate_result.mutate_millis
5959

6060
return LouvainMutateResult(**computation_result)
6161

@@ -156,7 +156,6 @@ def write(
156156
consecutive_ids: Optional[bool] = None,
157157
relationship_weight_property: Optional[str] = None,
158158
write_concurrency: Optional[int] = None,
159-
write_to_result_store: Optional[bool] = None,
160159
min_community_size: Optional[int] = None,
161160
) -> LouvainWriteResult:
162161
config = self._node_property_endpoints.create_base_config(
@@ -175,7 +174,7 @@ def write(
175174
seed_property=seed_property,
176175
sudo=sudo,
177176
tolerance=tolerance,
178-
write_to_result_store=write_to_result_store,
177+
write_concurrency=write_concurrency,
179178
)
180179

181180
computation_result, write_millis = self._node_property_endpoints.run_job_and_write(

‎graphdatascience/procedure_surface/arrow/node_property_endpoints.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33

44
from pandas import DataFrame
55

6+
from graphdatascience.arrow_client.v2.api_types import MutateResult
7+
68
from ...arrow_client.authenticated_flight_client import AuthenticatedArrowClient
79
from ...arrow_client.v2.data_mapper_utils import deserialize_single
810
from ...arrow_client.v2.job_client import JobClient
@@ -30,12 +32,12 @@ def run_job_and_get_summary(self, endpoint: str, G: Graph, config: Dict[str, Any
3032

3133
def run_job_and_mutate(
3234
self, endpoint: str, G: Graph, config: Dict[str, Any], mutate_property: str
33-
) -> tuple[Dict[str, Any], int]:
35+
) -> tuple[Dict[str, Any], MutateResult]:
3436
"""Run a job, mutate node properties, and return summary with mutation result."""
3537
job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config)
3638
mutate_result = MutationClient.mutate_node_property(self._arrow_client, job_id, mutate_property)
3739
computation_result = JobClient.get_summary(self._arrow_client, job_id)
38-
return computation_result, mutate_result.node_properties_written
40+
return computation_result, mutate_result
3941

4042
def run_job_and_stream(self, endpoint: str, G: Graph, config: Dict[str, Any]) -> DataFrame:
4143
"""Run a job and return streamed results."""

‎graphdatascience/procedure_surface/arrow/scc_arrow_endpoints.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,12 @@ def mutate(
3838
sudo=sudo,
3939
)
4040

41-
computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
41+
computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
4242
"v2/community.scc", G, config, mutate_property
4343
)
4444

45-
computation_result["nodePropertiesWritten"] = node_properties_written
46-
computation_result["mutateMillis"] = 0
45+
computation_result["nodePropertiesWritten"] = mutate_result.node_properties_written
46+
computation_result["mutateMillis"] = mutate_result.mutate_millis
4747

4848
return SccMutateResult(**computation_result)
4949

@@ -112,7 +112,6 @@ def write(
112112
job_id: Optional[str] = None,
113113
consecutive_ids: Optional[bool] = None,
114114
write_concurrency: Optional[int] = None,
115-
write_to_result_store: Optional[bool] = None,
116115
) -> SccWriteResult:
117116
config = self._node_property_endpoints.create_base_config(
118117
G,

0 commit comments

Comments
 (0)