[MCC-1476014] Python Lib-Create Fetch_data() function

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:

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: ...

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:

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())

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()

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."""

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

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]