Skip to content

Commit

Permalink
fix: Preseve dtypes modules in from dict (#1458)
Browse files Browse the repository at this point in the history
* fix: preserve stableness of dtypes in `from-dict`

* fixup
  • Loading branch information
MarcoGorelli authored Nov 28, 2024
1 parent 14a5f84 commit c08a09d
Show file tree
Hide file tree
Showing 13 changed files with 97 additions and 8 deletions.
5 changes: 5 additions & 0 deletions narwhals/_arrow/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,11 @@ def __narwhals_dataframe__(self) -> Self:
def __narwhals_lazyframe__(self) -> Self:
return self

def _change_dtypes(self, dtypes: DTypes) -> Self:
return self.__class__(
self._native_frame, backend_version=self._backend_version, dtypes=dtypes
)

def _from_native_frame(self, df: pa.Table) -> Self:
return self.__class__(
df, backend_version=self._backend_version, dtypes=self._dtypes
Expand Down
8 changes: 8 additions & 0 deletions narwhals/_arrow/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,14 @@ def __init__(
self._backend_version = backend_version
self._dtypes = dtypes

def _change_dtypes(self, dtypes: DTypes) -> Self:
return self.__class__(
self._native_series,
name=self._name,
backend_version=self._backend_version,
dtypes=dtypes,
)

def _from_native_series(self, series: pa.ChunkedArray | pa.Array) -> Self:
import pyarrow as pa # ignore-banned-import()

Expand Down
5 changes: 5 additions & 0 deletions narwhals/_dask/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,11 @@ def __narwhals_namespace__(self) -> DaskNamespace:
def __narwhals_lazyframe__(self) -> Self:
return self

def _change_dtypes(self, dtypes: DTypes) -> Self:
return self.__class__(
self._native_frame, backend_version=self._backend_version, dtypes=dtypes
)

def _from_native_frame(self, df: Any) -> Self:
return self.__class__(
df, backend_version=self._backend_version, dtypes=self._dtypes
Expand Down
3 changes: 3 additions & 0 deletions narwhals/_duckdb/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,5 +139,8 @@ def to_pandas(self: Self) -> pd.DataFrame:
def to_arrow(self: Self) -> pa.Table:
return self._native_frame.arrow()

def _change_dtypes(self: Self, dtypes: DTypes) -> Self:
return self.__class__(self._native_frame, dtypes=dtypes)

def _from_native_frame(self: Self, df: Any) -> Self:
return self.__class__(df, dtypes=self._dtypes)
3 changes: 3 additions & 0 deletions narwhals/_ibis/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,5 +119,8 @@ def __getattr__(self, attr: str) -> Any:
)
raise NotImplementedError(msg)

def _change_dtypes(self: Self, dtypes: DTypes) -> Self:
return self.__class__(self._native_frame, dtypes=dtypes)

def _from_native_frame(self: Self, df: Any) -> Self:
return self.__class__(df, dtypes=self._dtypes)
13 changes: 13 additions & 0 deletions narwhals/_interchange/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,11 +73,24 @@ def map_interchange_dtype_to_narwhals_dtype(
raise AssertionError(msg)


class WrapInterchangeFrame:
def __init__(self, interchange_frame: InterchangeFrame) -> None:
self._interchange_frame = interchange_frame

def __dataframe__(self) -> InterchangeFrame:
return self._interchange_frame


class InterchangeFrame:
def __init__(self, df: Any, dtypes: DTypes) -> None:
self._interchange_frame = df.__dataframe__()
self._dtypes = dtypes

def _change_dtypes(self: Self, dtypes: DTypes) -> Self:
return self.__class__(
WrapInterchangeFrame(self._interchange_frame), dtypes=dtypes
)

def __narwhals_dataframe__(self) -> Any:
return self

Expand Down
8 changes: 8 additions & 0 deletions narwhals/_pandas_like/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,14 @@ def _validate_columns(self, columns: pd.Index) -> None:
msg = f"Expected unique column names, got:{msg}"
raise ValueError(msg)

def _change_dtypes(self, dtypes: DTypes) -> Self:
return self.__class__(
self._native_frame,
implementation=self._implementation,
backend_version=self._backend_version,
dtypes=dtypes,
)

def _from_native_frame(self, df: Any) -> Self:
return self.__class__(
df,
Expand Down
8 changes: 8 additions & 0 deletions narwhals/_pandas_like/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,14 @@ def __getitem__(self, idx: int | slice | Sequence[int]) -> Any | Self:
return self._native_series.iloc[idx]
return self._from_native_series(self._native_series.iloc[idx])

def _change_dtypes(self, dtypes: DTypes) -> Self:
return self.__class__(
self._native_series,
implementation=self._implementation,
backend_version=self._backend_version,
dtypes=dtypes,
)

def _from_native_series(self, series: Any) -> Self:
return self.__class__(
series,
Expand Down
10 changes: 10 additions & 0 deletions narwhals/_polars/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,11 @@ def __native_namespace__(self: Self) -> ModuleType:
msg = f"Expected polars, got: {type(self._implementation)}" # pragma: no cover
raise AssertionError(msg)

def _change_dtypes(self, dtypes: DTypes) -> Self:
return self.__class__(
self._native_frame, backend_version=self._backend_version, dtypes=dtypes
)

def _from_native_frame(self, df: Any) -> Self:
return self.__class__(
df, backend_version=self._backend_version, dtypes=self._dtypes
Expand Down Expand Up @@ -321,6 +326,11 @@ def _from_native_frame(self, df: Any) -> Self:
df, backend_version=self._backend_version, dtypes=self._dtypes
)

def _change_dtypes(self, dtypes: DTypes) -> Self:
return self.__class__(
self._native_frame, backend_version=self._backend_version, dtypes=dtypes
)

def __getattr__(self, attr: str) -> Any:
def func(*args: Any, **kwargs: Any) -> Any:
import polars as pl # ignore-banned-import
Expand Down
5 changes: 5 additions & 0 deletions narwhals/_polars/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,11 @@ def __native_namespace__(self: Self) -> ModuleType:
msg = f"Expected polars, got: {type(self._implementation)}" # pragma: no cover
raise AssertionError(msg)

def _change_dtypes(self, dtypes: DTypes) -> Self:
return self.__class__(
self._native_series, backend_version=self._backend_version, dtypes=dtypes
)

def _from_native_series(self, series: Any) -> Self:
return self.__class__(
series, backend_version=self._backend_version, dtypes=self._dtypes
Expand Down
8 changes: 5 additions & 3 deletions narwhals/stable/v1/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -1160,19 +1160,21 @@ def _stableify(obj: Any) -> Any: ...
def _stableify(
obj: NwDataFrame[IntoFrameT] | NwLazyFrame[IntoFrameT] | NwSeries | NwExpr | Any,
) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series | Expr | Any:
from narwhals.stable.v1 import dtypes

if isinstance(obj, NwDataFrame):
return DataFrame(
obj._compliant_frame,
obj._compliant_frame._change_dtypes(dtypes),
level=obj._level,
)
if isinstance(obj, NwLazyFrame):
return LazyFrame(
obj._compliant_frame,
obj._compliant_frame._change_dtypes(dtypes),
level=obj._level,
)
if isinstance(obj, NwSeries):
return Series(
obj._compliant_series,
obj._compliant_series._change_dtypes(dtypes),
level=obj._level,
)
if isinstance(obj, NwExpr):
Expand Down
16 changes: 11 additions & 5 deletions tests/from_dict_test.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from __future__ import annotations

from datetime import datetime

import pytest

import narwhals as nw
Expand Down Expand Up @@ -61,12 +63,16 @@ def test_from_dict_one_native_one_narwhals(
def test_from_dict_v1(constructor: Constructor, request: pytest.FixtureRequest) -> None:
if "dask" in str(constructor):
request.applymarker(pytest.mark.xfail)
df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]}))
native_namespace = nw.get_native_namespace(df)
result = nw.from_dict({"c": [1, 2], "d": [5, 6]}, native_namespace=native_namespace)
expected = {"c": [1, 2], "d": [5, 6]}
df = nw_v1.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]}))
native_namespace = nw_v1.get_native_namespace(df)
result = nw_v1.from_dict(
{"c": [1, 2], "d": [datetime(2020, 1, 1), datetime(2020, 1, 2)]},
native_namespace=native_namespace,
)
expected = {"c": [1, 2], "d": [datetime(2020, 1, 1), datetime(2020, 1, 2)]}
assert_equal_data(result, expected)
assert isinstance(result, nw.DataFrame)
assert isinstance(result, nw_v1.DataFrame)
assert isinstance(result.schema["d"], nw_v1.dtypes.Datetime)


def test_from_dict_empty() -> None:
Expand Down
13 changes: 13 additions & 0 deletions tests/translate/from_native_test.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

from contextlib import nullcontext as does_not_raise
from typing import TYPE_CHECKING
from typing import Any

import numpy as np
Expand All @@ -13,6 +14,9 @@
import narwhals.stable.v1 as nw
from tests.utils import maybe_get_modin_df

if TYPE_CHECKING:
from narwhals.typing import DTypes

data = {"a": [1, 2, 3]}

df_pd = pd.DataFrame(data)
Expand All @@ -28,16 +32,25 @@


class MockDataFrame:
def _change_dtypes(self, _dtypes: DTypes) -> MockDataFrame:
return self

def __narwhals_dataframe__(self) -> Any:
return self


class MockLazyFrame:
def _change_dtypes(self, _dtypes: DTypes) -> MockLazyFrame:
return self

def __narwhals_lazyframe__(self) -> Any:
return self


class MockSeries:
def _change_dtypes(self, _dtypes: DTypes) -> MockSeries:
return self

def __narwhals_series__(self) -> Any:
return self

Expand Down

0 comments on commit c08a09d

Please sign in to comment.