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", ), ),