graphdatascience 1.21a2__tar.gz → 2.0a1__tar.gz

{graphdatascience-1.21a2/src/graphdatascience.egg-info → graphdatascience-2.0a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.21a2
+Version: 2.0a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Author-email: Neo4j <team-gds@neo4j.org>
 License-Expression: Apache-2.0
@@ -27,9 +27,9 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: multimethod<3.0,>=1.0
-Requires-Dist: neo4j<7.0,>=4.4.12
-Requires-Dist: numpy<2.4
-Requires-Dist: pandas<3.0,>=1.0
+Requires-Dist: neo4j<7.0,>=5.26.0
+Requires-Dist: numpy<3.0
+Requires-Dist: pandas<4.0,>=2.0
 Requires-Dist: pyarrow<24.0,>=18.0
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
@@ -42,7 +42,7 @@ Requires-Dist: networkx<4.0,>=2.0; extra == "networkx"
 Provides-Extra: ogb
 Requires-Dist: ogb<2.0,>=1.0; extra == "ogb"
 Provides-Extra: rust-ext
-Requires-Dist: neo4j-rust-ext<7.0,>=4.4.12; extra == "rust-ext"
+Requires-Dist: neo4j-rust-ext<7.0,>=5.26.0; extra == "rust-ext"
 Dynamic: license-file
 
 # Neo4j Graph Data Science Client

{graphdatascience-1.21a2 → graphdatascience-2.0a1}/pyproject.toml

@@ -29,9 +29,9 @@ classifiers = [
 requires-python = ">=3.10"
 dependencies = [
     "multimethod >= 1.0, < 3.0",
-    "neo4j >= 4.4.12, < 7.0",
-    "numpy < 2.4",
-    "pandas >= 1.0, < 3.0",
+    "neo4j >= 5.26.0, < 7.0",
+    "numpy < 3.0",
+    "pandas>=2.0,<4.0",
     "pyarrow >= 18.0, < 24.0",
     "textdistance >= 4.0, < 5.0",
     "tqdm >= 4.0, < 5.0",
@@ -50,7 +50,7 @@ BugTracker = "https://github.com/neo4j/graph-data-science-client/issues"
 [project.optional-dependencies]
 networkx = ["networkx >= 2.0, < 4.0"]
 ogb = ["ogb >= 1.0, < 2.0"]
-rust-ext = ["neo4j-rust-ext >= 4.4.12, < 7.0"]
+rust-ext = ["neo4j-rust-ext >= 5.26.0, < 7.0"]
 
 
 [dependency-groups]
@@ -96,22 +96,39 @@ notebook-base = [
     "nbformat==5.10.4",
     "nbclient==0.10.4",
     "ipykernel==7.1.0",
+    "neo4j-viz"
 ]
 notebook-aura-ci = [
     { include-group = "notebook-base" },
-    "python-dotenv==1.2.1"
+    "python-dotenv==1.2.1",
+    "pyspark[sql]"
 ]
 notebook-ci = [
     { include-group = "notebook-base" },
     "scipy == 1.14.0",
     "torch==2.3.0",
-    "torch-scatter==2.1.1",
-    "torch-sparse==0.6.17",
+    "torch-scatter==2.1.2",
+    "torch-sparse==0.6.18",
     "torch-geometric>=2.5.0",
 ]
 
 publish = ["build", "twine"]
 
+# torch-scatter/torch-sparse have no sdist-friendly install: building from source
+# requires a working C++/torch toolchain and fails on many local machines (notably
+# macOS arm64). Instead we pull prebuilt wheels from the PyG wheel index, which ships
+# wheels for both Linux (torch-2.3.0+cpu) and macOS (universal2).
+[tool.uv.sources]
+torch-scatter = [{ index = "pytorch-geometric-cpu" }]
+torch-sparse = [{ index = "pytorch-geometric-cpu" }]
+
+[[tool.uv.index]]
+name = "pytorch-geometric-cpu"
+url = "https://data.pyg.org/whl/torch-2.3.0+cpu.html"
+explicit = true
+format = "flat"
+
+# Fallback only: if uv ever has to build these from source, make torch available.
 [tool.uv.extra-build-dependencies]
 torch-sparse = ["torch"]
 torch-scatter = ["torch"]

graphdatascience-2.0a1/src/graphdatascience/__init__.py

@@ -0,0 +1,11 @@
+from .graph_data_science import GraphDataScience
+from .server_version.server_version import ServerVersion
+from .session import GdsSessions
+from .version import __version__
+
+__all__ = [
+    "__version__",
+    "ServerVersion",
+    "GraphDataScience",
+    "GdsSessions",
+]

{graphdatascience-1.21a2 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/arrow_info.py

@@ -4,7 +4,6 @@ from dataclasses import dataclass
 
 from graphdatascience.query_runner.query_runner import QueryRunner
 from graphdatascience.query_runner.query_type import QueryType
-from graphdatascience.server_version.server_version import ServerVersion
 
 
 @dataclass(frozen=True)
@@ -16,12 +15,11 @@ class ArrowInfo:
 
     @staticmethod
     def create(query_runner: QueryRunner) -> ArrowInfo:
-        debugYields = ["listenAddress", "enabled", "running"]
-        if query_runner.server_version() > ServerVersion(2, 6, 0):
-            debugYields.append("versions")
-
         procResult = query_runner.call_procedure(
-            endpoint="gds.debug.arrow", query_type=QueryType.SYSTEM, custom_error=False, yields=debugYields
+            endpoint="gds.debug.arrow",
+            query_type=QueryType.SYSTEM,
+            custom_error=False,
+            yields=["listenAddress", "enabled", "running", "versions"],
         ).iloc[0]
 
         return ArrowInfo(

graphdatascience-2.0a1/src/graphdatascience/arrow_client/arrow_table_utils.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import pandas
+import pyarrow
+from pyarrow import types as pa_types
+
+
+def _downcast_type(data_type: pyarrow.DataType) -> pyarrow.DataType:
+    """Map pyarrow "large" string/binary types back to their 32-bit variants.
+
+    pandas 3.0 introduced a default string dtype that ``pyarrow.Table.from_pandas``
+    maps to ``large_string`` (LargeUTF8), whereas pandas 2.x produced ``string``
+    (UTF8). The GDS server expects UTF8, so we coerce the large variants back.
+    List containers are preserved as-is, only their element type is downcast.
+    """
+    if pa_types.is_large_string(data_type):
+        return pyarrow.string()
+
+    return data_type
+
+
+def table_from_pandas(data: pandas.DataFrame) -> pyarrow.Table:
+    """Convert a DataFrame to a pyarrow Table, coercing large string/binary types to UTF8/binary.
+
+    Use this instead of ``pyarrow.Table.from_pandas`` so the produced schema stays
+    consistent across pandas 2.x and 3.0 (see :func:`_downcast_type`).
+    """
+    table = pyarrow.Table.from_pandas(data)
+
+    new_types = [_downcast_type(field.type) for field in table.schema]
+    if all(new_type == field.type for new_type, field in zip(new_types, table.schema)):
+        return table
+
+    new_schema = pyarrow.schema(
+        [field.with_type(new_type) for field, new_type in zip(table.schema, new_types)],
+        metadata=table.schema.metadata,
+    )
+    return table.cast(new_schema)

{graphdatascience-1.21a2 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/authenticated_flight_client.py

@@ -10,7 +10,7 @@ from typing import Any, Iterator, Type
 import certifi
 from pyarrow import Schema, flight
 from pyarrow import __version__ as arrow_version
-from pyarrow._flight import (
+from pyarrow.flight import (
     Action,
     ActionType,
     FlightInternalError,
@@ -22,7 +22,6 @@ from pyarrow._flight import (
 )
 
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
-from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from graphdatascience.retry_utils.retry_config import ExponentialWaitConfig, RetryConfigV2, StopConfig
 
 from ..version import __version__
@@ -32,90 +31,64 @@ from .middleware.user_agent_middleware import UserAgentFactory
 
 
 class AuthenticatedArrowClient:
-    @staticmethod
-    def create(
-        arrow_info: ArrowInfo,
-        auth: ArrowAuthentication | None = None,
-        encrypted: bool = False,
-        arrow_client_options: dict[str, Any] | None = None,
-        connection_string_override: str | None = None,
-        retry_config: RetryConfigV2 | None = None,
-        advertised_listen_address: tuple[str, int] | None = None,
-    ) -> AuthenticatedArrowClient:
-        connection_string: str
-        if connection_string_override is not None:
-            connection_string = connection_string_override
-        else:
-            connection_string = arrow_info.listenAddress
-
-        host, port = connection_string.split(":")
-
-        if retry_config is None:
-            retry_config = RetryConfigV2(
-                retryable_exceptions=[
-                    FlightTimedOutError,
-                    FlightUnavailableError,
-                    FlightInternalError,
-                ],
-                stop_config=StopConfig(after_delay=10, after_attempt=5),
-                wait_config=ExponentialWaitConfig(multiplier=1, min=1, max=10),
-            )
-
-        return AuthenticatedArrowClient(
-            host=host,
-            retry_config=retry_config,
-            port=int(port),
-            auth=auth,
-            encrypted=encrypted,
-            arrow_client_options=arrow_client_options,
-            advertised_listen_address=advertised_listen_address,
-        )
-
     def __init__(
         self,
-        host: str,
-        retry_config: RetryConfigV2,
-        port: int = 8491,
+        connection_info: str | tuple[str, int],
         auth: ArrowAuthentication | None = None,
         encrypted: bool = False,
         arrow_client_options: dict[str, Any] | None = None,
         user_agent: str | None = None,
         advertised_listen_address: tuple[str, int] | None = None,
+        retry_config: RetryConfigV2 | None = None,
     ):
         """Creates a new AuthenticatedArrowClient instance.
 
         Parameters
         ----------
-        host: str
-            The host address of the GDS Arrow server
-        port: int
-            The host port of the GDS Arrow server (default is 8491)
-        auth: ArrowAuthentication | None
+        connection_info
+            The host address and port of the GDS Arrow server
+        auth
             An implementation of ArrowAuthentication providing a pair to be used for basic authentication
-        encrypted: bool
+        encrypted
             A flag that indicates whether the connection should be encrypted (default is False)
-        arrow_client_options: dict[str, Any] | None
+        arrow_client_options
             Additional options to be passed to the Arrow Flight client.
-        user_agent: str | None
+        user_agent
             The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
-        retry_config: RetryConfig | None
+        retry_config
             The retry configuration to use for the Arrow requests send by the client.
-        advertised_listen_address: tuple[str, int] | None
+        advertised_listen_address
             The advertised listen address of the GDS Arrow server. This will be used by remote projection and writeback operations.
         """
+
+        if isinstance(connection_info, str):
+            host, port_str = connection_info.split(":")
+            port = int(port_str)
+        else:
+            host, port = connection_info
+
+        if retry_config is None:
+            retry_config = RetryConfigV2(
+                retryable_exceptions=[
+                    FlightTimedOutError,
+                    FlightUnavailableError,
+                    FlightInternalError,
+                ],
+                stop_config=StopConfig(after_delay=10, after_attempt=5),
+                wait_config=ExponentialWaitConfig(multiplier=1, min=1, max=10),
+            )
+
         self._host = host
-        self._port = port
-        self._auth = None
+        self._port = int(port)
+        self._auth = auth
         self._encrypted = encrypted
         self._arrow_client_options = arrow_client_options
         self._user_agent = user_agent
         self._logger = logging.getLogger("gds_arrow_client")
         self._retry_config = retry_config
         if auth:
-            self._auth = auth
             self._auth_middleware = AuthMiddleware(auth)
         self.advertised_listen_address = advertised_listen_address
-
         self._flight_client: flight.FlightClient = self._instantiate_flight_client()
 
     def connection_info(self) -> ConnectionInfo:
@@ -156,10 +129,14 @@ class AuthenticatedArrowClient:
 
         @self._retry_config.decorator(operation_name="Request token", logger=self._logger)
         def auth_with_retry() -> None:
-            client = self._flight_client
-            if self._auth:
-                auth_pair = self._auth.auth_pair()
-                client.authenticate_basic_token(auth_pair[0], auth_pair[1])
+            try:
+                client = self._flight_client
+                if self._auth:
+                    auth_pair = self._auth.auth_pair()
+                    client.authenticate_basic_token(auth_pair[0], auth_pair[1])
+            except (FlightTimedOutError, FlightUnavailableError, FlightInternalError):
+                self._reconnect()
+                raise
 
         if self._auth:
             auth_with_retry()
@@ -178,9 +155,13 @@ class AuthenticatedArrowClient:
     def do_action_with_retry(self, endpoint: str, payload: bytes | dict[str, Any]) -> list[Result]:
         @self._retry_config.decorator(operation_name="Send action", logger=self._logger)
         def run_with_retry() -> list[Result]:
-            # the Flight response error code is only checked on iterator consumption
-            # we eagerly collect iterator here to trigger retry in case of an error
-            return list(self.do_action(endpoint, payload))
+            try:
+                # the Flight response error code is only checked on iterator consumption
+                # we eagerly collect iterator here to trigger retry in case of an error
+                return list(self.do_action(endpoint, payload))
+            except (FlightTimedOutError, FlightUnavailableError, FlightInternalError):
+                self._reconnect()
+                raise
 
         return run_with_retry()
 
@@ -192,7 +173,11 @@ class AuthenticatedArrowClient:
     ) -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
         @self._retry_config.decorator(operation_name="Do put", logger=self._logger)
         def run_with_retry() -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
-            return self._flight_client.do_put(descriptor, schema)  # type: ignore
+            try:
+                return self._flight_client.do_put(descriptor, schema)  # type: ignore
+            except (FlightTimedOutError, FlightUnavailableError, FlightInternalError):
+                self._reconnect()
+                raise
 
         return run_with_retry()
 
@@ -247,6 +232,14 @@ class AuthenticatedArrowClient:
         self.__dict__.update(state)
         self._flight_client = self._instantiate_flight_client()
 
+    def _reconnect(self) -> None:
+        try:
+            self._flight_client.close()
+        except Exception:
+            pass
+
+        self._flight_client = self._instantiate_flight_client()
+
 
 @dataclass
 class ConnectionInfo:

{graphdatascience-1.21a2 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/middleware/auth_middleware.py

@@ -4,7 +4,7 @@ import base64
 import time
 from typing import Any
 
-from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
+from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
 
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
 

{graphdatascience-1.21a2 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/middleware/user_agent_middleware.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import Any
 
-from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
+from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
 
 
 class UserAgentFactory(ClientMiddlewareFactory):  # type: ignore

{graphdatascience-1.21a2/src/graphdatascience/arrow_client/v2 → graphdatascience-2.0a1/src/graphdatascience/arrow_client/v1}/data_mapper_utils.py

@@ -1,7 +1,7 @@
 import json
 from typing import Any
 
-from pyarrow._flight import Result
+from pyarrow.flight import Result
 
 
 def deserialize_single(input_stream: list[Result]) -> dict[str, Any]:

{graphdatascience-1.21a2 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/v1/gds_arrow_client.py

@@ -12,11 +12,11 @@ from pyarrow.types import is_dictionary
 from pydantic import BaseModel
 
 from graphdatascience.arrow_client.arrow_endpoint_version import ArrowEndpointVersion
+from graphdatascience.arrow_client.arrow_table_utils import table_from_pandas
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient, ConnectionInfo
 from graphdatascience.arrow_client.v1.data_mapper_utils import deserialize_single
 
 from ...procedure_surface.arrow.error_handler import handle_flight_error
-from ...semantic_version.semantic_version import SemanticVersion
 from ..progress_callback import ProgressCallback
 
 
@@ -487,11 +487,13 @@ class GdsArrowClient:
         batch_size: int,
         progress_callback: ProgressCallback,
     ) -> None:
+        batches: list[RecordBatch]
         match data:
             case pyarrow.Table():
                 batches = data.to_batches(batch_size)
             case pandas.DataFrame():
-                batches = pyarrow.Table.from_pandas(data).to_batches(batch_size)
+                batches = table_from_pandas(data).to_batches(batch_size)
+
             case _:
                 batches = data
 
@@ -562,11 +564,7 @@ class GdsArrowClient:
         except Exception as e:
             handle_flight_error(e)
         arrow_table = self._sanitize_arrow_table(arrow_table)
-        if SemanticVersion.from_string(pandas.__version__) >= SemanticVersion(2, 0, 0):
-            return arrow_table.to_pandas(types_mapper=pandas.ArrowDtype)  # type: ignore
-        else:
-            arrow_table = self._sanitize_arrow_table(arrow_table)
-            return arrow_table.to_pandas()  # type: ignore
+        return arrow_table.to_pandas(types_mapper=pandas.ArrowDtype)  # type: ignore
 
     def __enter__(self) -> GdsArrowClient:
         return self

{graphdatascience-1.21a2/src/graphdatascience/arrow_client/v1 → graphdatascience-2.0a1/src/graphdatascience/arrow_client/v2}/data_mapper_utils.py

@@ -1,7 +1,7 @@
 import json
 from typing import Any
 
-from pyarrow._flight import Result
+from pyarrow.flight import Result
 
 
 def deserialize_single(input_stream: list[Result]) -> dict[str, Any]:

{graphdatascience-1.21a2 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/v2/gds_arrow_client.py

@@ -10,6 +10,7 @@ import pyarrow
 from pyarrow import RecordBatch, flight
 
 from graphdatascience.arrow_client.arrow_endpoint_version import ArrowEndpointVersion
+from graphdatascience.arrow_client.arrow_table_utils import table_from_pandas
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient, ConnectionInfo
 from graphdatascience.query_runner.termination_flag import TerminationFlag
 
@@ -123,7 +124,8 @@ class GdsArrowClient:
             job_id=job_id,
         )
 
-        return JobClient.run_job_and_wait(self._flight_client, "v2/graph.nodeLabel.stream", config, log_progress)
+        compute_job = JobClient.run_job_and_wait(self._flight_client, "v2/graph.nodeLabel.stream", config, log_progress)
+        return JobClient.start_export_result(self._flight_client, graph_name, compute_job)
 
     def get_relationships(
         self,
@@ -175,14 +177,12 @@ class GdsArrowClient:
 
         return JobClient.run_job(self._flight_client, endpoint, config)
 
-    def stream_job(self, graph_name: str, job_id: str) -> pandas.DataFrame:
+    def stream_job(self, job_id: str) -> pandas.DataFrame:
         """
         Streams the results of a previously started job.
 
         Parameters
         ----------
-        graph_name
-           The graph name.
         job_id
             Identifier for the computation.
 
@@ -191,7 +191,7 @@ class GdsArrowClient:
         DataFrame
             A pandas DataFrame containing the results of the job.
         """
-        return JobClient().stream_results(self._flight_client, graph_name, job_id)
+        return JobClient().get_stream(self._flight_client, job_id)
 
     def create_graph(
         self,
@@ -415,7 +415,7 @@ class GdsArrowClient:
         job_id : str
             The job id of the process
         """
-        self._flight_client.do_action_with_retry("v2/jobs.cancel", {"jobId": job_id})
+        JobClient.cancel_job(self._flight_client, job_id)
 
     def job_status(self, job_id: str) -> JobStatus:
         """
@@ -472,32 +472,21 @@ class GdsArrowClient:
 
         return self._flight_client.request_token()
 
+    def flight_client(self) -> AuthenticatedArrowClient:
+        """
+        Returns the underlying flight client.
+        """
+        return self._flight_client
+
     def _upload_data(
         self,
         endpoint: str,
         job_id: str,
-        data: pyarrow.Table | list[pyarrow.RecordBatch] | list[pandas.DataFrame],
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | list[pandas.DataFrame] | pandas.DataFrame,
         batch_size: int = 10000,
         progress_callback: ProgressCallback = lambda num_rows: None,
         termination_flag: TerminationFlag | None = None,
     ) -> None:
-        if isinstance(data, pandas.DataFrame):
-            data = [data]
-
-        def map_list_entry_to_batches(e: pyarrow.RecordBatch | pandas.DataFrame) -> list[pyarrow.RecordBatch]:
-            return [e] if isinstance(e, pyarrow.RecordBatch) else pyarrow.Table.from_pandas(e).to_batches(batch_size)
-
-        batches: list[pyarrow.RecordBatch] = []
-        match data:
-            case list():
-                for entry in data:
-                    partial_batches = map_list_entry_to_batches(entry)
-                    batches.extend(partial_batches)
-            case pyarrow.Table():
-                batches = data.to_batches(batch_size)
-            case _:
-                raise ValueError(f"Unsupported data type: {type(data)}")
-
         flight_descriptor = {
             "name": endpoint,
             "version": ArrowEndpointVersion.V2.version(),
@@ -507,11 +496,45 @@ class GdsArrowClient:
         }
         upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
 
+        for batches in self._batch_groups(data, batch_size):
+            self._upload_batches(upload_descriptor, batches, job_id, progress_callback, termination_flag)
+
+    @staticmethod
+    def _batch_groups(
+        data: pyarrow.Table | list[pyarrow.RecordBatch] | list[pandas.DataFrame] | pandas.DataFrame,
+        batch_size: int,
+    ) -> list[list[pyarrow.RecordBatch]]:
+        if isinstance(data, pandas.DataFrame):
+            return [table_from_pandas(data).to_batches(batch_size)]
+
+        if isinstance(data, pyarrow.Table):
+            return [data.to_batches(batch_size)]
+
+        if isinstance(data, list):
+            if all(isinstance(entry, pandas.DataFrame) for entry in data):
+                return [table_from_pandas(entry).to_batches(batch_size) for entry in data]
+
+            if all(isinstance(entry, pyarrow.RecordBatch) for entry in data):
+                return [data]
+
+        raise ValueError(f"Unsupported data type: {type(data)}")
+
+    def _upload_batches(
+        self,
+        upload_descriptor: flight.FlightDescriptor,
+        batches: list[pyarrow.RecordBatch],
+        job_id: str,
+        progress_callback: ProgressCallback,
+        termination_flag: TerminationFlag | None,
+    ) -> None:
+        if not batches:
+            return
+
         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)
+        def upload_batch(batch: RecordBatch) -> None:
+            put_stream.write_batch(batch)
 
         with put_stream:
             for partition in batches:

{graphdatascience-1.21a2 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/v2/job_client.py

@@ -2,17 +2,18 @@ import json
 from typing import Any
 
 from pandas import ArrowDtype, DataFrame
-from pyarrow._flight import Ticket
+from pyarrow.flight import Ticket
 from tenacity import Retrying, retry_if_result
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.api_types import JobIdConfig, JobStatus
 from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
-from graphdatascience.query_runner.progress.progress_bar import TqdmProgressBar
+from graphdatascience.progress.progress_bar import TqdmProgressBar
 from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.retry_utils.retry_utils import job_wait_strategy
 
 JOB_STATUS_ENDPOINT = "v2/jobs.status"
+JOBS_CANCEL_ENDPOINT = "v2/jobs.cancel"
 RESULTS_SUMMARY_ENDPOINT = "v2/results.summary"
 
 
@@ -73,6 +74,10 @@ class JobClient:
                     if progress_bar:
                         progress_bar.update(job_status.status, job_status.progress_percent(), job_status.sub_tasks())
 
+    @staticmethod
+    def cancel_job(client: AuthenticatedArrowClient, job_id: str) -> None:
+        client.do_action_with_retry(JOBS_CANCEL_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
+
     @staticmethod
     def get_job_status(client: AuthenticatedArrowClient, job_id: str) -> JobStatus:
         arrow_res = client.do_action_with_retry(JOB_STATUS_ENDPOINT, JobIdConfig(jobId=job_id).dump_camel())
@@ -86,14 +91,22 @@ class JobClient:
 
     @staticmethod
     def stream_results(client: AuthenticatedArrowClient, graph_name: str, job_id: str) -> DataFrame:
+        export_job_id = JobClient.start_export_result(client, graph_name, job_id)
+
+        return JobClient.get_stream(client, export_job_id)
+
+    @staticmethod
+    def start_export_result(client: AuthenticatedArrowClient, graph_name: str, job_id: str) -> str:
         payload = {
             "graphName": graph_name,
             "jobId": job_id,
         }
 
         res = client.do_action_with_retry("v2/results.stream", payload)
-        export_job_id = JobIdConfig(**deserialize_single(res)).job_id
+        return JobIdConfig(**deserialize_single(res)).job_id
 
+    @staticmethod
+    def get_stream(client: AuthenticatedArrowClient, export_job_id: str) -> DataFrame:
         stream_payload = {"version": "v2", "name": export_job_id, "body": {}}
 
         ticket = Ticket(json.dumps(stream_payload).encode("utf-8"))