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

dataconnect/client.py

@@ -6,7 +6,7 @@
 from types import TracebackType
 from typing import Any
 
-import pyarrow as pa
+import pandas as pd
 
 from dataconnect import _encoding
 from dataconnect.auth import BearerTokenAuth
@@ -60,13 +60,40 @@ def datasets(self, study_uuid: str) -> list[Dataset]:
         rows = self._action_json(_ACTION_LIST_DATASETS, {"study_uuid": body})
         return [Dataset(**r) for r in rows]
 
-    def fetch_data(self, dataset_uuid: str) -> pa.Table:
-        """Fetch the data for a given dataset as a PyArrow Table."""
-        body = {"dataset_uuid": dataset_uuid}
-        results = self._transport.do_action(_ACTION_FETCH_TICKET, _encoding.dumps(body))
-        if not results:
-            raise RuntimeError("Server returned no data for the fetch_data action.")
-        return self._transport.do_get(results).read_all()
+    def fetch_data(
+        self,
+        dataset_uuid: str,
+        first_n_rows: int | None = None,
+    ) -> pd.DataFrame:
+        """Fetch data of a single dataset as a pandas DataFrame.
+
+        Streams Arrow record batches over Flight and assembles them into a
+        single pandas DataFrame, preserving clinical data types
+        (dates, decimals, etc.).
+
+        Parameters
+        ----------
+        dataset_uuid:
+            UUID of the target dataset (required). The Study and Study
+            Environment contexts are resolved automatically server-side.
+        first_n_rows:
+            Optional row limit. When provided, only the first
+            ``first_n_rows`` rows of the dataset are returned. When omitted
+            (``None``), the full dataset is returned.
+
+        Returns
+        -------
+        pandas.DataFrame
+            The dataset materialized as a pandas DataFrame.
+        """
+        if first_n_rows is not None and first_n_rows <= 0:
+            raise ValueError("first_n_rows must be a positive integer when provided.")
+
+        return _get_dataset(
+            transport=self._transport,
+            dataset_uuid=dataset_uuid,
+            limit=first_n_rows,
+        )
 
     # Lifecycle
     def close(self) -> None:
@@ -91,3 +118,30 @@ def _action_json(self, action: str, body: dict[str, Any] | None) -> Any:
         if not results:
             return []
         return json.loads(results.decode("utf-8"))
+
+
+def _get_dataset(
+    transport: FlightTransport,
+    dataset_uuid: str,
+    limit: int | None,
+) -> pd.DataFrame:
+    """Fetch a single dataset over Arrow Flight as a pandas DataFrame.
+
+    Mirrors the R implementation in ``R/datasets.R::.get_dataset`` /
+    ``.get_dataset_raw``: builds a JSON Flight ticket and streams record
+    batches into a single :class:`pandas.DataFrame`.
+    """
+    if not dataset_uuid or not str(dataset_uuid).strip():
+        raise ValueError("dataset_uuid must be provided.")
+
+    ticket_data: dict[str, Any] = {
+        "dataset_uuid": dataset_uuid,
+        "dataset_name": "",
+        "limit": limit,
+    }
+
+    ticket = _encoding.dumps(ticket_data)
+    stream = transport.do_get(ticket)
+
+    table = stream.read_all()
+    return table.to_pandas()

dataconnect/framework/pyarrow_transport.py

@@ -8,7 +8,7 @@
 from pyarrow import flight
 
 from dataconnect.auth import BearerTokenAuth, Credentials
-from dataconnect.exceptions import AuthenticationError, DataConnectError, QueryError
+from dataconnect.exceptions import AuthenticationError, ConnectionError, DataConnectError, QueryError
 from dataconnect.framework.transport import FlightTransport, RecordBatchStream
 
 

pyproject.toml

@@ -27,13 +27,6 @@ optional = true
 [tool.poetry.group.service.dependencies]
 gunicorn ="^20.1.0"
 
-# ML
-[tool.poetry.group.ml]
-optional = true
-
-[tool.poetry.group.ml.dependencies]
-pandas = "^2.0.2"
-
 # DEV
 [tool.poetry.group.dev]
 optional = true

tests/test_client.py

@@ -1,5 +1,90 @@
+"""Tests for DataConnectClient, including fetch_data and _get_dataset."""
+
+from __future__ import annotations
+
+import datetime as dt
+import json
+from collections.abc import Iterator
+from decimal import Decimal
+
+import pandas as pd
+import pyarrow as pa
 import pytest
 
+from dataconnect import _encoding
+from dataconnect.client import DataConnectClient, _get_dataset
+from dataconnect.framework.transport import FlightTransport, RecordBatchStream
+
+
+# ---------------------------------------------------------------------------
+# Fakes
+# ---------------------------------------------------------------------------
+
+class _FakeStream(RecordBatchStream):
+    def __init__(self, batches: list[pa.RecordBatch]) -> None:
+        self._batches = batches
+
+    def read_all(self) -> pa.Table:
+        if not self._batches:
+            return pa.table({})
+        return pa.Table.from_batches(self._batches)
+
+    def __iter__(self) -> Iterator[pa.RecordBatch]:
+        yield from self._batches
+
+
+class _FakeTransport(FlightTransport):
+    def __init__(self, batches: list[pa.RecordBatch]) -> None:
+        self._batches = batches
+        self.last_ticket: bytes | None = None
+        self.do_get_calls = 0
+
+    def do_action(self, action: str, body: bytes = b"") -> bytes:
+        return b""
+
+    def do_get(self, ticket: bytes) -> RecordBatchStream:
+        self.last_ticket = ticket
+        self.do_get_calls += 1
+        return _FakeStream(self._batches)
+
+    def do_put(self, command: bytes, table: pa.Table) -> bytes | None:
+        return b""
+
+    def close(self) -> None:
+        return None
+
+
+def _make_clinical_batches() -> list[pa.RecordBatch]:
+    """Two batches with date32 and decimal128 columns to exercise dtype roundtrip."""
+    schema = pa.schema(
+        [
+            pa.field("subject_id", pa.string()),
+            pa.field("visit_date", pa.date32()),
+            pa.field("dose_mg", pa.decimal128(precision=10, scale=2)),
+        ]
+    )
+    batch_a = pa.RecordBatch.from_pydict(
+        {
+            "subject_id": ["S1", "S2"],
+            "visit_date": [dt.date(2026, 1, 15), dt.date(2026, 2, 1)],
+            "dose_mg": [Decimal("12.50"), Decimal("7.25")],
+        },
+        schema=schema,
+    )
+    batch_b = pa.RecordBatch.from_pydict(
+        {
+            "subject_id": ["S3"],
+            "visit_date": [dt.date(2026, 3, 10)],
+            "dose_mg": [Decimal("100.00")],
+        },
+        schema=schema,
+    )
+    return [batch_a, batch_b]
+
+
+# ---------------------------------------------------------------------------
+# Placeholder / benchmark
+# ---------------------------------------------------------------------------
 
 def test_client() -> None:
     assert True
@@ -9,3 +94,106 @@ def test_client() -> None:
 def test_dummy_benchmark() -> None:
     # Dummy benchmark test to satisfy CI
     assert True
+
+
+# ---------------------------------------------------------------------------
+# fetch_data
+# ---------------------------------------------------------------------------
+
+def test_fetch_data_returns_pandas_dataframe_preserving_types() -> None:
+    transport = _FakeTransport(_make_clinical_batches())
+    client = DataConnectClient(transport)
+
+    df = client.fetch_data(dataset_uuid="ds-123")
+
+    assert isinstance(df, pd.DataFrame)
+    assert list(df.columns) == ["subject_id", "visit_date", "dose_mg"]
+    assert len(df) == 3
+    # date32 → python date objects (preserved, not coerced to ns timestamps)
+    assert df["visit_date"].iloc[0] == dt.date(2026, 1, 15)
+    # decimal128 → Decimal objects (precision preserved)
+    assert df["dose_mg"].iloc[0] == Decimal("12.50")
+    assert df["dose_mg"].iloc[2] == Decimal("100.00")
+
+
+def test_fetch_data_sends_ticket_with_dataset_uuid_and_limit() -> None:
+    transport = _FakeTransport(_make_clinical_batches())
+    client = DataConnectClient(transport)
+
+    client.fetch_data(dataset_uuid="ds-123", first_n_rows=10)
+
+    assert transport.last_ticket is not None
+    payload = json.loads(transport.last_ticket.decode("utf-8"))
+    assert payload["dataset_uuid"] == "ds-123"
+    assert payload["limit"] == 10
+    assert payload["dataset_name"] == ""
+    # Study / study-env UUIDs are no longer part of the ticket.
+    assert "study_uuid" not in payload
+    assert "study_env_uuid" not in payload
+
+
+def test_fetch_data_omitted_first_n_rows_sends_null_limit() -> None:
+    transport = _FakeTransport(_make_clinical_batches())
+    client = DataConnectClient(transport)
+
+    client.fetch_data(dataset_uuid="ds-123")
+
+    payload = json.loads(transport.last_ticket.decode("utf-8"))  # type: ignore[union-attr]
+    assert payload["limit"] is None
+
+
+def test_fetch_data_invalid_first_n_rows_raises() -> None:
+    transport = _FakeTransport(_make_clinical_batches())
+    client = DataConnectClient(transport)
+
+    with pytest.raises(ValueError):
+        client.fetch_data(dataset_uuid="ds-123", first_n_rows=0)
+
+
+def test_fetch_data_empty_stream_returns_empty_dataframe() -> None:
+    transport = _FakeTransport([])
+    client = DataConnectClient(transport)
+
+    df = client.fetch_data(dataset_uuid="ds-123")
+
+    assert isinstance(df, pd.DataFrame)
+    assert df.empty
+
+
+def test_fetch_data_rejects_legacy_study_kwargs() -> None:
+    """The legacy study_uuid / study_environment_uuid kwargs must be removed."""
+    transport = _FakeTransport(_make_clinical_batches())
+    client = DataConnectClient(transport)
+
+    with pytest.raises(TypeError):
+        client.fetch_data(dataset_uuid="ds-123", study_uuid="study-1")  # type: ignore[call-arg]
+    with pytest.raises(TypeError):
+        client.fetch_data(dataset_uuid="ds-123", study_environment_uuid="env-1")  # type: ignore[call-arg]
+
+
+# ---------------------------------------------------------------------------
+# _get_dataset helper
+# ---------------------------------------------------------------------------
+
+def test_get_dataset_helper_builds_expected_ticket() -> None:
+    transport = _FakeTransport(_make_clinical_batches())
+
+    _get_dataset(
+        transport=transport,
+        dataset_uuid="ds-9",
+        limit=5,
+    )
+
+    payload = _encoding.loads(transport.last_ticket)  # type: ignore[arg-type]
+    assert payload == {
+        "dataset_uuid": "ds-9",
+        "dataset_name": "",
+        "limit": 5,
+    }
+
+
+def test_get_dataset_helper_requires_dataset_uuid() -> None:
+    transport = _FakeTransport(_make_clinical_batches())
+
+    with pytest.raises(ValueError):
+        _get_dataset(transport=transport, dataset_uuid="", limit=None)