diff --git a/dataconnect/client.py b/dataconnect/client.py index 9a6f270..2476201 100644 --- a/dataconnect/client.py +++ b/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() diff --git a/dataconnect/framework/pyarrow_transport.py b/dataconnect/framework/pyarrow_transport.py index df7d49a..fb07497 100644 --- a/dataconnect/framework/pyarrow_transport.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index f8732d8..4b0348e 100644 --- a/pyproject.toml +++ b/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 diff --git a/tests/test_client.py b/tests/test_client.py index cb45e43..c54338a 100644 --- a/tests/test_client.py +++ b/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)