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

README.md

@@ -17,8 +17,7 @@ dependency.
 ## Installation
 
 ```bash
-pip install dataconnect          # core (pyarrow + pydantic + httpx)
-pip install dataconnect[pandas]  # + pandas for .to_pandas() on results
+pip install dataconnect          # core (pyarrow + pandas)
 ```
 
 Requires **Python ≥ 3.13**.

dataconnect/client.py

@@ -6,12 +6,13 @@
 from types import TracebackType
 from typing import Any
 
+import pandas as pd
 import pyarrow as pa
 
 from dataconnect import _encoding
 from dataconnect.auth import BearerTokenAuth
 from dataconnect.framework.pyarrow_transport import PyArrowFlightTransport
-from dataconnect.framework.transport import FlightTransport
+from dataconnect.framework.transport import FlightTransport, RecordBatchStream
 from dataconnect.models import Dataset, Study
 
 # Flight actions / commands
@@ -60,13 +61,76 @@ 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:
+    def fetch_data(self, dataset_uuid: str, first_n_rows: int | None = None) -> pd.DataFrame:
+        """Fetch data for a dataset and return a pandas DataFrame.
+
+        Data is transferred using Arrow Flight DoGet and converted to pandas in
+        streaming batches to reduce peak memory overhead.
+        """
+        if not dataset_uuid or not dataset_uuid.strip():
+            raise ValueError("dataset_uuid must be a non-empty string.")
+
+        if first_n_rows is not None:
+            if isinstance(first_n_rows, bool) or not isinstance(first_n_rows, int):
+                raise TypeError("first_n_rows must be an integer when provided.")
+            if first_n_rows < 0:
+                raise ValueError("first_n_rows must be >= 0 when provided.")
+
+        ticket_payload: dict[str, Any] = {
+            "study_uuid": None,
+            "study_env_uuid": None,
+            "dataset_uuid": dataset_uuid,
+            "dataset_name": "",
+        }
+
+        if first_n_rows is not None:
+            ticket_payload["limit"] = int(first_n_rows)
+
+        ticket = self._transport.do_action(_ACTION_FETCH_TICKET, _encoding.dumps(ticket_payload))
+        if not ticket:
             raise RuntimeError("Server returned no data for the fetch_data action.")
-        return self._transport.do_get(results).read_all()
+
+        stream = self._transport.do_get(ticket)
+        return self._stream_to_pandas(stream, first_n_rows)
+
+    def _stream_to_pandas(
+        self,
+        stream: RecordBatchStream,
+        first_n_rows: int | None,
+    ) -> pd.DataFrame:
+        """Convert a record-batch stream into a pandas DataFrame."""
+        if first_n_rows == 0:
+            return pd.DataFrame()
+
+        frames: list[pd.DataFrame] = []
+        remaining = first_n_rows
+
+        for batch in stream:
+            current_batch = batch
+            if remaining is not None:
+                if remaining <= 0:
+                    break
+                if batch.num_rows > remaining:
+                    current_batch = batch.slice(0, remaining)
+
+            frames.append(
+                current_batch.to_pandas(
+                    types_mapper=pd.ArrowDtype,
+                    date_as_object=False,
+                    timestamp_as_object=False,
+                )
+            )
+
+            if remaining is not None:
+                remaining -= current_batch.num_rows
+
+        if not frames:
+            return pd.DataFrame()
+
+        if len(frames) == 1:
+            return frames[0]
+
+        return pd.concat(frames, ignore_index=True, copy=False)
 
     # Lifecycle
     def close(self) -> None:

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]
@@ -27,13 +28,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