diff --git a/kalshi/models/common.py b/kalshi/models/common.py index f5cd3692..097c68ea 100644 --- a/kalshi/models/common.py +++ b/kalshi/models/common.py @@ -51,8 +51,9 @@ def to_dataframe(self) -> pandas.DataFrame: "Install it with: pip install 'kalshi-sdk[pandas]'" ) from exc - records = [item.model_dump(mode="python") for item in self.items] - return pd.DataFrame(records) + if not self.items: + return pd.DataFrame() + return pd.DataFrame(self._columns()) def to_polars(self) -> polars.DataFrame: """Return page items as a polars DataFrame (requires kalshi-sdk[polars]).""" @@ -64,5 +65,24 @@ def to_polars(self) -> polars.DataFrame: "Install it with: pip install 'kalshi-sdk[polars]'" ) from exc - records = [item.model_dump(mode="python") for item in self.items] - return pl.DataFrame(records) + if not self.items: + return pl.DataFrame() + return pl.DataFrame(self._columns()) + + def _columns(self) -> dict[str, list[object]]: + """Build a column-oriented dict[field, list[value]] over self.items. + + Avoids the per-row ``model_dump`` walk used previously (issue #264). + Nested ``BaseModel`` cells are dumped to dicts per column so that + polars can still infer a ``Struct`` dtype and the #225 / #190 + Decimal contract is preserved (DollarDecimal lands as ``Decimal`` + in ``mode="python"``, not as ``str``). + """ + fields = type(self.items[0]).model_fields.keys() + columns: dict[str, list[object]] = {} + for field in fields: + col: list[object] = [getattr(item, field) for item in self.items] + if col and isinstance(col[0], BaseModel): + col = [v.model_dump(mode="python") if isinstance(v, BaseModel) else v for v in col] + columns[field] = col + return columns diff --git a/tests/test_page_dataframe.py b/tests/test_page_dataframe.py index c7ab96fb..c8bb1937 100644 --- a/tests/test_page_dataframe.py +++ b/tests/test_page_dataframe.py @@ -55,9 +55,7 @@ def test_empty_page(self) -> None: assert isinstance(df, pd.DataFrame) assert df.empty - def test_missing_pandas_raises_importerror( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_missing_pandas_raises_importerror(self, monkeypatch: pytest.MonkeyPatch) -> None: # Ensure no cached pandas module satisfies the import. monkeypatch.delitem(sys.modules, "pandas", raising=False) real_import = builtins.__import__ @@ -105,9 +103,7 @@ def test_empty_page(self) -> None: assert isinstance(df, pl.DataFrame) assert df.is_empty() - def test_missing_polars_raises_importerror( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_missing_polars_raises_importerror(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delitem(sys.modules, "polars", raising=False) real_import = builtins.__import__ @@ -287,9 +283,7 @@ def test_page_to_dataframe_preserves_decimal_objects(self) -> None: def test_page_to_dataframe_decimal_sum_is_decimal_not_concatenated_strings(self) -> None: pytest.importorskip("pandas") - page: Page[_DollarRow] = Page( - items=_dollar_items("0.5600", "0.5600"), cursor=None - ) + page: Page[_DollarRow] = Page(items=_dollar_items("0.5600", "0.5600"), cursor=None) df = page.to_dataframe() total = df["price"].sum() @@ -311,9 +305,7 @@ def test_page_to_polars_preserves_decimal_objects(self) -> None: def test_page_to_polars_decimal_sum_is_decimal(self) -> None: pytest.importorskip("polars") - page: Page[_DollarRow] = Page( - items=_dollar_items("0.5600", "0.5600"), cursor=None - ) + page: Page[_DollarRow] = Page(items=_dollar_items("0.5600", "0.5600"), cursor=None) df = page.to_polars() total = df["price"].sum() @@ -321,3 +313,64 @@ def test_page_to_polars_decimal_sum_is_decimal(self) -> None: # polars Decimal.sum() returns a Python Decimal — never a string concat. assert isinstance(total, Decimal) assert total == Decimal("1.1200") + + +# --------------------------------------------------------------------------- +# Issue #264: column-oriented build must preserve the Decimal contract at +# scale and across both backends. Pins behavior for a 100-row page (the +# size at which per-row model_dump was measurably slower). +# --------------------------------------------------------------------------- + + +class TestPageToDataframeColumnOriented: + def test_hundred_row_decimal_column_holds_decimal_instances(self) -> None: + pd = pytest.importorskip("pandas") + items = _dollar_items(*[f"0.{i:04d}" for i in range(100)]) + page: Page[_DollarRow] = Page(items=items, cursor=None) + + df = page.to_dataframe() + + assert df.shape == (100, 1) + assert df["price"].dtype == object + values = df["price"].tolist() + assert all(isinstance(v, Decimal) for v in values) + assert not any(isinstance(v, float) for v in values) + assert values[0] == Decimal("0.0000") + assert values[-1] == Decimal("0.0099") + # sentinel: pandas attribute survives + assert isinstance(df, pd.DataFrame) + + def test_empty_page_returns_empty_dataframe(self) -> None: + pd = pytest.importorskip("pandas") + page: Page[_DollarRow] = Page(items=[], cursor=None) + + df = page.to_dataframe() + + assert isinstance(df, pd.DataFrame) + assert df.empty + assert df.shape == (0, 0) + + +class TestPageToPolarsColumnOriented: + def test_hundred_row_decimal_column_holds_decimal_instances(self) -> None: + pl = pytest.importorskip("polars") + items = _dollar_items(*[f"0.{i:04d}" for i in range(100)]) + page: Page[_DollarRow] = Page(items=items, cursor=None) + + df = page.to_polars() + + assert df.shape == (100, 1) + assert isinstance(df["price"].dtype, pl.Decimal) + values = df["price"].to_list() + assert all(isinstance(v, Decimal) for v in values) + assert values[0] == Decimal("0.0000") + assert values[-1] == Decimal("0.0099") + + def test_empty_page_returns_empty_polars_dataframe(self) -> None: + pl = pytest.importorskip("polars") + page: Page[_DollarRow] = Page(items=[], cursor=None) + + df = page.to_polars() + + assert isinstance(df, pl.DataFrame) + assert df.shape == (0, 0)