From 583f329e28a96ba95f9da61374153d935aa78378 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Mon, 4 Aug 2025 14:44:35 +0000 Subject: [PATCH] feat: switch to pyarrow types for consistency --- leanframe/core/frame.py | 9 ++++++++- leanframe/core/series.py | 4 +++- tests/unit/test_series.py | 7 +++++-- 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/leanframe/core/frame.py b/leanframe/core/frame.py index 9173014..505fb71 100644 --- a/leanframe/core/frame.py +++ b/leanframe/core/frame.py @@ -50,4 +50,11 @@ def __getitem__(self, key: str): return leanframe.core.series.Series(self._data[key]) def to_pandas(self) -> pandas.DataFrame: - return self._data.to_pandas() + """Convert the DataFrame to a pandas.DataFrame. + + Where possible, pandas.ArrowDtype is used to avoid lossy conversions + from the database types to pandas. + """ + return self._data.to_pyarrow().to_pandas( + types_mapper=lambda type_: pandas.ArrowDtype(type_) + ) diff --git a/leanframe/core/series.py b/leanframe/core/series.py index fc4c0b1..837292a 100644 --- a/leanframe/core/series.py +++ b/leanframe/core/series.py @@ -37,4 +37,6 @@ def name(self) -> str: def to_pandas(self) -> pandas.Series: """Convert to a pandas Series.""" - return self._data.to_pandas() + return self._data.to_pyarrow().to_pandas( + types_mapper=lambda type_: pandas.ArrowDtype(type_) + ) diff --git a/tests/unit/test_series.py b/tests/unit/test_series.py index 61ee8d5..1f3eb20 100644 --- a/tests/unit/test_series.py +++ b/tests/unit/test_series.py @@ -16,6 +16,7 @@ import pandas import pandas.testing +import pyarrow import pytest import leanframe @@ -25,11 +26,13 @@ ("series_pd",), ( pytest.param( - pandas.Series([1, 2, 3]), + pandas.Series([1, 2, 3], dtype=pandas.ArrowDtype(pyarrow.int64())), id="int64", ), pytest.param( - pandas.Series([1.0, float("nan"), 3.0]), + pandas.Series( + [1.0, float("nan"), 3.0], dtype=pandas.ArrowDtype(pyarrow.float64()) + ), id="float64", ), ),