diff --git a/dataconnect/client.py b/dataconnect/client.py index 84ce20d..83a0258 100644 --- a/dataconnect/client.py +++ b/dataconnect/client.py @@ -8,6 +8,9 @@ from __future__ import annotations from types import TracebackType +from uuid import UUID + +import pandas as pd from dataconnect.models import Study from dataconnect.service import DataConnectService, DefaultDataConnectService @@ -46,6 +49,14 @@ def get_studies(self) -> list[Study]: """List the studies the client is authorized to access.""" return self._service.get_studies() + def fetch_data( + self, + dataset_uuid: UUID, + first_n_rows: int | None = None, + ) -> pd.DataFrame: + """Fetch data frames for a given dataset UUID.""" + return self._service.fetch_data(dataset_uuid, first_n_rows) + # Lifecycle def close(self) -> None: diff --git a/dataconnect/service/base.py b/dataconnect/service/base.py index 78a26b0..da54bab 100644 --- a/dataconnect/service/base.py +++ b/dataconnect/service/base.py @@ -3,6 +3,9 @@ from __future__ import annotations from abc import ABC, abstractmethod +from uuid import UUID + +import pandas as pd from dataconnect.models import Study @@ -13,5 +16,12 @@ class DataConnectService(ABC): @abstractmethod def get_studies(self) -> list[Study]: ... + @abstractmethod + def fetch_data( + self, + dataset_uuid: UUID, + first_n_rows: int | None = None, + ) -> pd.DataFrame: ... + @abstractmethod def close(self) -> None: ... diff --git a/dataconnect/service/default.py b/dataconnect/service/default.py index dd287cf..8a1fc7b 100644 --- a/dataconnect/service/default.py +++ b/dataconnect/service/default.py @@ -2,6 +2,10 @@ from __future__ import annotations +from uuid import UUID + +import pandas as pd + from dataconnect.exceptions import ( AuthenticationError, AuthorizationError, @@ -14,7 +18,7 @@ ) from dataconnect.models import Study from dataconnect.service.base import DataConnectService -from dataconnect.service.mappers import resource_to_study +from dataconnect.service.mappers import resource_to_study, resource_to_fetched_data from dataconnect.transport.base import Transport from dataconnect.transport.errors import ( TransportAuthenticationError, @@ -29,6 +33,7 @@ # Server action identifiers _ACTION_LIST_STUDIES = "studies.list" +_ACTION_FETCH_TICKET = "data.fetch_ticket" def _translate_error(ex: TransportError) -> DataConnectError: @@ -72,6 +77,29 @@ def get_studies(self) -> list[Study]: except (IndexError, KeyError, TypeError, ValueError) as ex: raise ValidationError(f"Unexpected studies response format: {ex}") from ex + def fetch_data(self, dataset_uuid: UUID, first_n_rows: int | None = None) -> pd.DataFrame: + + if not dataset_uuid or not str(dataset_uuid).strip(): + raise ValueError("dataset_uuid must be provided.") + + if first_n_rows is not None and first_n_rows <= 0: + raise ValueError("first_n_rows must be a positive integer when provided.") + + request = ResourceQuery(action=_ACTION_FETCH_TICKET).append_body( + { + "study_env_uuid": None, + "dataset_name": None, + "dataset_uuid": str(dataset_uuid), + "limit": first_n_rows, + } + ) + + try: + table = self._transport.do_get(request) + return resource_to_fetched_data(table) + except TransportError as ex: + raise _translate_error(ex) from ex + def close(self) -> None: try: diff --git a/dataconnect/service/mappers.py b/dataconnect/service/mappers.py index a4a6407..807cd9e 100644 --- a/dataconnect/service/mappers.py +++ b/dataconnect/service/mappers.py @@ -10,9 +10,12 @@ import json from uuid import UUID +import pandas as pd +import pyarrow as pa + from dataconnect.exceptions import NotFoundError from dataconnect.models import Study, StudyEnvironment -from dataconnect.transport.models import ResourceInfo +from dataconnect.transport.models import ResourceInfo, DataTable def resource_to_study(resource: ResourceInfo) -> Study: @@ -28,3 +31,13 @@ def resource_to_study(resource: ResourceInfo) -> Study: name=data["name"], environments=[StudyEnvironment(uuid=UUID(e["uuid"]), name=e["name"]) for e in data.get("environments", [])], ) + + +def resource_to_fetched_data(table: DataTable) -> pd.DataFrame: + """Convert a transport-layer ``DataTable`` into a ``pandas.DataFrame``.""" + + ipc_buffer = pa.BufferReader(table.ipc_bytes) + reader = pa.ipc.open_stream(ipc_buffer) + table = reader.read_all() + + return pd.DataFrame(table.to_pandas()) \ No newline at end of file diff --git a/dataconnect/transport/arrow_flight/transport.py b/dataconnect/transport/arrow_flight/transport.py index cd714ee..680e39f 100644 --- a/dataconnect/transport/arrow_flight/transport.py +++ b/dataconnect/transport/arrow_flight/transport.py @@ -9,6 +9,7 @@ import json +import pyarrow as pa from pyarrow import flight from dataconnect.transport.base import Transport @@ -18,7 +19,7 @@ TransportConnectionError, TransportStatusError, ) -from dataconnect.transport.models import DataRef, ResourceInfo, ResourceQuery +from dataconnect.transport.models import DataRef, ResourceInfo, ResourceQuery, DataTable def _to_resource_info(info: flight.FlightInfo) -> ResourceInfo: @@ -35,10 +36,28 @@ def _to_resource_info(info: flight.FlightInfo) -> ResourceInfo: total_records=info.total_records, ) +def _to_bytes(table: pa.Table) -> DataTable: + """Serialize a ``pa.Table`` to a technology-agnostic ``DataTable``. + + Each record batch is serialized individually as Arrow IPC bytes. + The schema is serialized separately so it can be recovered without + the data batches. + """ + schema_bytes = table.schema.serialize().to_pybytes() + + sink = pa.BufferOutputStream() + writer = pa.ipc.new_stream(sink, table.schema) + for batch in table.to_batches(): + writer.write_batch(batch) + writer.close() + ipc_bytes = sink.getvalue().to_pybytes() + + return DataTable(schema_bytes=schema_bytes, ipc_bytes=ipc_bytes) # Maps service-layer action names to the flight_type value the Arrow Flight server expects. _ACTION_FLIGHT_TYPE: dict[str, str] = { "studies.list": "STUDIES", + "data.fetch_ticket": "DATA_FETCH_TICKET", } @@ -101,5 +120,46 @@ def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: except Exception as ex: raise TransportConnectionError(f"Unexpected error during list_resources: {ex}") from ex + def do_get(self, request: ResourceQuery) -> DataTable: + """Call FlightClient.do_get and read all chunks into a single pa.Table.""" + + flight_type = _ACTION_FLIGHT_TYPE.get(request.action) + + if flight_type is None: + raise TransportStatusError( + f"Unknown action: {request.action!r}", status_code=3, grpc_status="INVALID_ARGUMENT" + ) + + body = json.loads(request.body) if request.body else {} + ticket_bytes = json.dumps(body, separators=(",", ":")).encode("utf-8") + ticket = flight.Ticket(ticket_bytes) + + try: + table = self._client.do_get(ticket, self._options()) + batches: list[pa.RecordBatch] = [] + while True: + try: + chunk, _metadata = table.read_chunk() + batches.append(chunk) + except StopIteration: + break + except flight.FlightError as ex: + raise TransportConnectionError(f"Error reading stream: {ex}") from ex + + return _to_bytes(pa.Table.from_batches(batches)) # validate schema + batches can be serialized + + except flight.FlightUnauthenticatedError as ex: + raise TransportAuthenticationError(str(ex)) from ex + except flight.FlightUnauthorizedError as ex: + raise TransportAuthorizationError(str(ex)) from ex + except flight.FlightUnavailableError as ex: + raise TransportConnectionError(str(ex)) from ex + except flight.FlightInternalError as ex: + raise TransportStatusError(str(ex), status_code=13, grpc_status="INTERNAL") from ex + except flight.FlightError as ex: + raise TransportConnectionError(str(ex)) from ex + except Exception as ex: + raise TransportConnectionError(f"Unexpected error during do_get: {ex}") from ex + def close(self) -> None: self._client.close() diff --git a/dataconnect/transport/base.py b/dataconnect/transport/base.py index df12c50..fe8b757 100644 --- a/dataconnect/transport/base.py +++ b/dataconnect/transport/base.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod -from dataconnect.transport.models import ResourceInfo, ResourceQuery +from dataconnect.transport.models import ResourceInfo, ResourceQuery, DataTable class Transport(ABC): @@ -23,6 +23,14 @@ def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: service layer's responsibility. """ + @abstractmethod + def do_get(self, request: ResourceQuery) -> DataTable: + """Retrieve data for a single endpoint ticket. + + Reads all chunks from the server stream and returns a single + ``DataTable`` containing the complete result set. + """ + @abstractmethod def close(self) -> None: """Close the transport connection.""" diff --git a/dataconnect/transport/models.py b/dataconnect/transport/models.py index 65ea576..b7074b4 100644 --- a/dataconnect/transport/models.py +++ b/dataconnect/transport/models.py @@ -38,3 +38,13 @@ class ResourceInfo: endpoints: list[DataRef] total_records: int schema_bytes: bytes + +@dataclass(frozen=True) +class DataTable: + """Technology-agnostic representation of a fetched data result. + + ``schema_bytes`` holds the Arrow IPC-serialized schema. + ``ipc_bytes`` holds the full Arrow IPC stream (schema + all batches). + """ + schema_bytes: bytes + ipc_bytes: bytes \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index a0d5680..61f165d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ url = "https://mdsol.jfrog.io/artifactory/api/pypi/pypi-prod-virtual/simple" [tool.poetry.dependencies] python = "^3.13" pyarrow = "^19.0.0" +pandas = "^2.0.2" # SERVICE [tool.poetry.group.service] @@ -32,7 +33,6 @@ gunicorn ="^20.1.0" optional = true [tool.poetry.group.ml.dependencies] -pandas = "^2.0.2" # DEV [tool.poetry.group.dev]