Address PR comments

Author: DarthMax

graphdatascience/arrow_client/v2/gds_arrow_client.py

@@ -173,7 +173,7 @@ def get_relationships(
 
         return JobClient.run_job(self._flight_client, endpoint, config)
 
-    def stream(self, graph_name: str, job_id: str) -> pandas.DataFrame:
+    def stream_job(self, graph_name: str, job_id: str) -> pandas.DataFrame:
         """
         Streams the results of a previously started job.
 
@@ -391,44 +391,7 @@ def upload_triplets(
         """
         self._upload_data("graph.project.fromTriplets", job_id, data, batch_size, progress_callback)
 
-    def _upload_data(
-        self,
-        endpoint: str,
-        job_id: str,
-        data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
-        batch_size: int = 10000,
-        progress_callback: Callable[[int], None] = lambda x: None,
-    ) -> None:
-        match data:
-            case pyarrow.Table():
-                batches = data.to_batches(batch_size)
-            case pandas.DataFrame():
-                batches = pyarrow.Table.from_pandas(data).to_batches(batch_size)
-            case _:
-                batches = data
-
-        flight_descriptor = {
-            "name": endpoint,
-            "version": ArrowEndpointVersion.V2.version(),
-            "body": {
-                "jobId": job_id,
-            },
-        }
-        upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
-
-        put_stream, ack_stream = self._flight_client.do_put_with_retry(upload_descriptor, batches[0].schema)
-
-        @self._flight_client._retry_config.decorator(operation_name="Upload batch", logger=self._logger)
-        def upload_batch(p: RecordBatch) -> None:
-            put_stream.write_batch(p)
-
-        with put_stream:
-            for partition in batches:
-                upload_batch(partition)
-                ack_stream.read()
-                progress_callback(partition.num_rows)
-
-    def abort(self, job_id: str) -> None:
+    def abort_job(self, job_id: str) -> None:
         """
         Aborts the specified process
 
@@ -494,6 +457,43 @@ def request_token(self) -> str | None:
 
         return self._flight_client.request_token()
 
+    def _upload_data(
+            self,
+            endpoint: str,
+            job_id: str,
+            data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
+            batch_size: int = 10000,
+            progress_callback: Callable[[int], None] = lambda x: None,
+    ) -> None:
+        match data:
+            case pyarrow.Table():
+                batches = data.to_batches(batch_size)
+            case pandas.DataFrame():
+                batches = pyarrow.Table.from_pandas(data).to_batches(batch_size)
+            case _:
+                batches = data
+
+        flight_descriptor = {
+            "name": endpoint,
+            "version": ArrowEndpointVersion.V2.version(),
+            "body": {
+                "jobId": job_id,
+            },
+        }
+        upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
+
+        put_stream, ack_stream = self._flight_client.do_put_with_retry(upload_descriptor, batches[0].schema)
+
+        @self._flight_client._retry_config.decorator(operation_name="Upload batch", logger=self._logger)
+        def upload_batch(p: RecordBatch) -> None:
+            put_stream.write_batch(p)
+
+        with put_stream:
+            for partition in batches:
+                upload_batch(partition)
+                ack_stream.read()
+                progress_callback(partition.num_rows)
+
     def __enter__(self) -> GdsArrowClient:
         return self
 

graphdatascience/tests/integrationV2/arrow_client/v2/test_gds_arrow_client.py

@@ -50,27 +50,28 @@ def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[GraphV2, N
 
 def test_stream_node_label(gds_arrow_client: GdsArrowClient, sample_graph: GraphV2) -> None:
     job_id = gds_arrow_client.get_nodes(sample_graph.name(), node_filter="n.prop1 > 1")
-    result = gds_arrow_client.stream(sample_graph.name(), job_id)
+    result = gds_arrow_client.stream_job(sample_graph.name(), job_id)
 
     assert ["nodeId"] == list(result.columns)
     assert len(result) == 2
 
 
 def test_stream_node_properties(gds_arrow_client: GdsArrowClient, sample_graph: GraphV2) -> None:
     job_id = gds_arrow_client.get_node_properties(sample_graph.name(), node_properties=["prop1", "prop2"])
-    result = gds_arrow_client.stream(sample_graph.name(), job_id)
+    result = gds_arrow_client.stream_job(sample_graph.name(), job_id)
 
     assert len(result) == 3
     assert "nodeId" in result.columns
     assert "prop1" in result.columns
     assert "prop2" in result.columns
+    assert {"nodeId", "prop1", "prop2"} == set(result.columns)
     assert set(result["prop1"].tolist()) == {1, 2, 3}
     assert set(result["prop2"].tolist()) == {42.0, 43.0, 44.0}
 
 
 def test_stream_relationship_properties(gds_arrow_client: GdsArrowClient, sample_graph: GraphV2) -> None:
     job_id = gds_arrow_client.get_relationships(sample_graph.name(), ["REL"], relationship_properties=["relX", "relY"])
-    result = gds_arrow_client.stream(sample_graph.name(), job_id)
+    result = gds_arrow_client.stream_job(sample_graph.name(), job_id)
 
     assert len(result) == 2
     assert "sourceNodeId" in result.columns