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

dataconnect/client.py

@@ -10,6 +10,8 @@
 from types import TracebackType
 from uuid import UUID
 
+import pandas as pd
+
 from dataconnect.models import Dataset, DatasetVersion, Study
 from dataconnect.service import DataConnectService, DefaultDataConnectService
 
@@ -51,6 +53,14 @@ def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]:
         """List the dataset versions the client is authorized to access."""
         return self._service.get_dataset_versions(dataset_uuid)
 
+    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)
+
     def get_datasets(
         self,
         study_environment_uuid: UUID,

dataconnect/service/base.py

@@ -5,6 +5,8 @@
 from abc import ABC, abstractmethod
 from uuid import UUID
 
+import pandas as pd
+
 from dataconnect.models import Dataset, DatasetVersion, Study
 
 
@@ -26,5 +28,12 @@ def get_datasets(
     @abstractmethod
     def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]: ...
 
+    @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

@@ -4,6 +4,8 @@
 
 from uuid import UUID
 
+import pandas as pd
+
 from dataconnect.exceptions import (
     AuthenticationError,
     AuthorizationError,
@@ -16,7 +18,12 @@
 )
 from dataconnect.models import Dataset, DatasetVersion, Study
 from dataconnect.service.base import DataConnectService
-from dataconnect.service.mappers import resource_to_dataset, resource_to_dataset_version, resource_to_study
+from dataconnect.service.mappers import (
+    resource_to_dataset,
+    resource_to_dataset_version,
+    resource_to_fetched_data,
+    resource_to_study,
+)
 from dataconnect.service.validators import validate_search_study_name
 from dataconnect.transport.base import Transport
 from dataconnect.transport.errors import (
@@ -34,6 +41,7 @@
 _ACTION_LIST_STUDIES = "studies.list"
 _ACTION_LIST_DATASETS = "datasets.list"
 _ACTION_LIST_DATASET_VERSIONS = "dataset_versions.list"
+_ACTION_FETCH_TICKET = "data.fetch_ticket"
 
 
 def _translate_error(ex: TransportError) -> DataConnectError:
@@ -101,6 +109,32 @@ def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]:
         except (IndexError, KeyError, TypeError, ValueError) as ex:
             raise ValidationError(f"Unexpected dataset versions 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 dataset_uuid.int == 0:
+            raise ValueError("dataset_uuid must not be an empty UUID.")
+
+        if first_n_rows is not None and (not isinstance(first_n_rows, int) or 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 get_datasets(
         self,
         study_environment_uuid: UUID,

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 Dataset, DatasetVersion, Study, StudyEnvironment
-from dataconnect.transport.models import ResourceInfo
+from dataconnect.transport.models import DataTable, ResourceInfo
 
 
 def resource_to_study(resource: ResourceInfo) -> Study:
@@ -47,6 +50,14 @@ def resource_to_dataset_version(resource: ResourceInfo) -> DatasetVersion:
     )
 
 
+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)
+    with pa.ipc.open_stream(ipc_buffer) as reader:
+        return reader.read_all().to_pandas()
+
+
 def resource_to_dataset(resource: ResourceInfo) -> Dataset:
     """Parse a transport-layer ``ResourceInfo`` into a ``Dataset`` domain object."""
 

dataconnect/transport/arrow_flight/transport.py

@@ -12,6 +12,7 @@
 import platform
 import subprocess
 
+import pyarrow as pa
 import pyarrow.flight as flight
 
 from dataconnect.transport.base import Transport
@@ -21,7 +22,7 @@
     TransportConnectionError,
     TransportStatusError,
 )
-from dataconnect.transport.models import DataRef, ResourceInfo, ResourceQuery
+from dataconnect.transport.models import DataRef, DataTable, ResourceInfo, ResourceQuery
 
 
 def _to_resource_info(info: flight.FlightInfo) -> ResourceInfo:
@@ -39,11 +40,31 @@ def _to_resource_info(info: flight.FlightInfo) -> ResourceInfo:
     )
 
 
+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",
     "datasets.list": "DATASETS",
     "dataset_versions.list": "VERSIONS",
+    "data.fetch_ticket": "DATA_FETCH_TICKET",
 }
 
 
@@ -141,5 +162,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, "flight_type": flight_type}, 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, schema=table.schema))
+
+        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 DataTable, ResourceInfo, ResourceQuery
 
 
 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:
+        """Fetch the full dataset described by ``request``.
+
+        All record batches from the stream are read and returned as a single
+        ``DataTable`` containing the complete result set.
+        """
+
     @abstractmethod
     def close(self) -> None:
         """Close the transport connection."""

dataconnect/transport/models.py

@@ -38,3 +38,15 @@ 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]