From 89d055f3dde8ebb3cd07ec1541b5b4715a8792f3 Mon Sep 17 00:00:00 2001 From: Alex Fieraru Date: Mon, 4 May 2026 16:15:19 +0300 Subject: [PATCH 1/2] feat: create fetch_data() function --- dataconnect/client.py | 71 ++++++++++++++++++++++++++++++++++++++----- 1 file changed, 64 insertions(+), 7 deletions(-) diff --git a/dataconnect/client.py b/dataconnect/client.py index 9a6f270..1ad9e3b 100644 --- a/dataconnect/client.py +++ b/dataconnect/client.py @@ -6,6 +6,7 @@ from types import TracebackType from typing import Any +import pandas as pd import pyarrow as pa from dataconnect import _encoding @@ -60,13 +61,69 @@ 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 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) + + stream = self._transport.do_get(_encoding.dumps(ticket_payload)) + return self._stream_to_pandas(stream, first_n_rows) + + def _stream_to_pandas( + self, + stream: pa.RecordBatchReader | Any, + 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() + + return pd.concat(frames, ignore_index=True, copy=False) # Lifecycle def close(self) -> None: From 412ffaf4e616d8e6b1f67622b9787d0b47f4c198 Mon Sep 17 00:00:00 2001 From: Alex Fieraru Date: Mon, 4 May 2026 16:29:40 +0300 Subject: [PATCH 2/2] feat: fix copilot recommendaitons --- README.md | 3 +-- dataconnect/client.py | 13 ++++++++++--- pyproject.toml | 8 +------- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 186a77f..33acda2 100644 --- a/README.md +++ b/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**. diff --git a/dataconnect/client.py b/dataconnect/client.py index 1ad9e3b..8bd415d 100644 --- a/dataconnect/client.py +++ b/dataconnect/client.py @@ -12,7 +12,7 @@ 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 @@ -86,12 +86,16 @@ def fetch_data(self, dataset_uuid: str, first_n_rows: int | None = None) -> pd.D if first_n_rows is not None: ticket_payload["limit"] = int(first_n_rows) - stream = self._transport.do_get(_encoding.dumps(ticket_payload)) + 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.") + + stream = self._transport.do_get(ticket) return self._stream_to_pandas(stream, first_n_rows) def _stream_to_pandas( self, - stream: pa.RecordBatchReader | Any, + stream: RecordBatchStream, first_n_rows: int | None, ) -> pd.DataFrame: """Convert a record-batch stream into a pandas DataFrame.""" @@ -123,6 +127,9 @@ def _stream_to_pandas( if not frames: return pd.DataFrame() + if len(frames) == 1: + return frames[0] + return pd.concat(frames, ignore_index=True, copy=False) # Lifecycle diff --git a/pyproject.toml b/pyproject.toml index f8732d8..c341671 100644 --- a/pyproject.toml +++ b/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