Repository: narwhals-dev/narwhals
Pull request #147 (2024-05-12 13:07:07)

add to_frame, maybe_align_index, get_native_namespace


Patch:
diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md
index 1d2a7cd0..8425ed17 100644
--- a/docs/api-reference/narwhals.md
+++ b/docs/api-reference/narwhals.md
@@ -10,7 +10,9 @@ Here are the top-level functions available in Narwhals.
         - col
         - concat
         - from_native
+        - get_native_namespace
         - len
+        - maybe_align_index
         - max
         - mean
         - min
diff --git a/docs/api-reference/series.md b/docs/api-reference/series.md
index fc35c1a9..e87a30f7 100644
--- a/docs/api-reference/series.md
+++ b/docs/api-reference/series.md
@@ -28,6 +28,7 @@
         - sort
         - std
         - sum
+        - to_frame
         - to_numpy
         - to_pandas
         - unique
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index 156b3bb7..d3038b49 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -25,7 +25,9 @@ from narwhals.expression import sum_horizontal
 from narwhals.functions import concat
 from narwhals.series import Series
 from narwhals.translate import from_native
+from narwhals.translate import get_native_namespace
 from narwhals.translate import to_native
+from narwhals.utils import maybe_align_index
 
 __version__ = "0.8.11"
 
@@ -33,6 +35,8 @@ __all__ = [
     "concat",
     "to_native",
     "from_native",
+    "maybe_align_index",
+    "get_native_namespace",
     "all",
     "col",
     "len",
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index d1398ddf..e00ae9dc 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -12,6 +12,9 @@ from narwhals._pandas_like.utils import horizontal_concat
 from narwhals._pandas_like.utils import translate_dtype
 from narwhals._pandas_like.utils import validate_dataframe_comparand
 from narwhals._pandas_like.utils import validate_indices
+from narwhals.translate import get_cudf
+from narwhals.translate import get_modin
+from narwhals.translate import get_pandas
 from narwhals.utils import flatten
 
 if TYPE_CHECKING:
@@ -49,6 +52,19 @@ class PandasDataFrame:
 
         return PandasNamespace(self._implementation)
 
+    def __native_namespace__(self) -> Any:
+        if self._implementation == "pandas":
+            return get_pandas()
+        if self._implementation == "modin":  # pragma: no cover
+            return get_modin()
+        if self._implementation == "cudf":  # pragma: no cover
+            return get_cudf()
+        msg = f"Expected pandas/modin/cudf, got: {type(self._implementation)}"  # pragma: no cover
+        raise AssertionError(msg)
+
+    def __len__(self) -> int:
+        return len(self._dataframe)
+
     def _validate_columns(self, columns: Sequence[str]) -> None:
         if len(columns) != len(set(columns)):
             counter = collections.Counter(columns)
@@ -104,7 +120,9 @@ class PandasDataFrame:
 
     def with_row_index(self, name: str) -> Self:
         row_index = create_native_series(
-            range(len(self._dataframe)), implementation=self._implementation
+            range(len(self._dataframe)),
+            index=self._dataframe.index,
+            implementation=self._implementation,
         ).alias(name)
         return self._from_dataframe(
             horizontal_concat(
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 897e7ffc..660dc8fa 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -139,6 +139,13 @@ class PandasSeries:
     def item(self) -> Any:
         return item(self._series)
 
+    def to_frame(self) -> Any:
+        from narwhals._pandas_like.dataframe import PandasDataFrame
+
+        return PandasDataFrame(
+            self._series.to_frame(), implementation=self._implementation
+        )
+
     def is_between(
         self, lower_bound: Any, upper_bound: Any, closed: str = "both"
     ) -> PandasSeries:
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index a818c120..325389d8 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -119,18 +119,22 @@ def parse_into_expr(implementation: str, into_expr: IntoPandasExpr) -> PandasExp
     raise AssertionError(msg)
 
 
-def create_native_series(iterable: Any, implementation: str) -> PandasSeries:
+def create_native_series(
+    iterable: Any,
+    implementation: str,
+    index: Any = None,
+) -> PandasSeries:
     from narwhals._pandas_like.series import PandasSeries
 
     if implementation == "pandas":
         pd = get_pandas()
-        series = pd.Series(iterable, name="")
+        series = pd.Series(iterable, index=index, name="")
     elif implementation == "modin":
         mpd = get_modin()
-        series = mpd.Series(iterable, name="")
+        series = mpd.Series(iterable, index=index, name="")
     elif implementation == "cudf":
         cudf = get_cudf()
-        series = cudf.Series(iterable, name="")
+        series = cudf.Series(iterable, index=index, name="")
     return PandasSeries(series, implementation=implementation)
 
 
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index fd1282c6..1d4f3777 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -31,11 +31,17 @@ class BaseFrame:
     _dataframe: Any
     _is_polars: bool
 
-    def __narwhals_namespace__(self) -> Any:
+    def __len__(self) -> Any:
+        return self._dataframe.__len__()
+
+    def __native_namespace__(self) -> Any:
         if self._is_polars:
-            import polars as pl
+            return get_polars()
+        return self._dataframe.__native_namespace__()
 
-            return pl
+    def __narwhals_namespace__(self) -> Any:
+        if self._is_polars:
+            return get_polars()
         return self._dataframe.__narwhals_namespace__()
 
     def _from_dataframe(self, df: Any) -> Self:
diff --git a/narwhals/series.py b/narwhals/series.py
index 2c883df2..f22083c8 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -14,6 +14,8 @@ if TYPE_CHECKING:
     import numpy as np
     from typing_extensions import Self
 
+    from narwhals.dataframe import DataFrame
+
 
 class Series:
     def __init__(
@@ -58,8 +60,10 @@ class Series:
     def __array__(self, *args: Any, **kwargs: Any) -> np.ndarray:
         return self._series.to_numpy(*args, **kwargs)
 
-    def __getitem__(self, idx: int) -> Any:
-        return self._series[idx]
+    def __getitem__(self, idx: int | slice) -> Any:
+        if isinstance(idx, int):
+            return self._series[idx]
+        return self._from_series(self._series[idx])
 
     def __narwhals_namespace__(self) -> Any:
         if self._is_polars:
@@ -163,6 +167,48 @@ class Series:
             self._series.cast(translate_dtype(self.__narwhals_namespace__(), dtype))
         )
 
+    def to_frame(self) -> DataFrame:
+        """
+        Convert to dataframe.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> s = [1, 2, 3]
+            >>> s_pd = pd.Series(s, name='a')
+            >>> s_pl = pl.Series('a', s)
+
+            We define a library agnostic function:
+
+            >>> def func(s_any):
+            ...     s = nw.from_native(s_any, series_only=True)
+            ...     df = s.to_frame()
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+               a
+            0  1
+            1  2
+            2  3
+            >>> func(s_pl)
+            shape: (3, 1)
+            ┌─────┐
+            │ a   │
+            │ --- │
+            │ i64 │
+            ╞═════╡
+            │ 1   │
+            │ 2   │
+            │ 3   │
+            └─────┘
+        """
+        from narwhals.dataframe import DataFrame
+
+        return DataFrame(self._series.to_frame())
+
     def mean(self) -> Any:
         """
         Reduce this Series to the mean value.
diff --git a/narwhals/translate.py b/narwhals/translate.py
index 984b3d52..9ce98932 100644
--- a/narwhals/translate.py
+++ b/narwhals/translate.py
@@ -205,10 +205,29 @@ def from_native(
     return native_dataframe  # type: ignore[no-any-return]  # pragma: no cover (todo)
 
 
+def get_native_namespace(obj: Any) -> Any:
+    """
+    Get native namespace from object.
+
+    Examples:
+        >>> import polars as pl
+        >>> import pandas as pd
+        >>> import narwhals as nw
+        >>> df = nw.from_native(pd.DataFrame({'a': [1,2,3]}))
+        >>> nw.get_native_namespace(df)
+        <module 'pandas'...>
+        >>> df = nw.from_native(pl.DataFrame({'a': [1,2,3]}))
+        >>> nw.get_native_namespace(df)
+        <module 'polars'...>
+    """
+    return obj.__native_namespace__()
+
+
 __all__ = [
     "get_pandas",
     "get_polars",
     "get_modin",
     "get_cudf",
+    "get_native_namespace",
     "to_native",
 ]
diff --git a/narwhals/utils.py b/narwhals/utils.py
index cf8afaf2..81b02279 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -1,14 +1,20 @@
 from __future__ import annotations
 
 import re
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterable
 from typing import Sequence
 from typing import TypeVar
+from typing import cast
 
 from narwhals.dependencies import get_pandas
 from narwhals.dependencies import get_polars
 
+if TYPE_CHECKING:
+    from narwhals.dataframe import BaseFrame
+    from narwhals.series import Series
+
 T = TypeVar("T")
 
 
@@ -73,3 +79,85 @@ def validate_same_library(items: Iterable[Any]) -> None:
     ):
         return
     raise NotImplementedError("Cross-library comparisons aren't supported")
+
+
+def maybe_align_index(lhs: T, rhs: Series | BaseFrame) -> T:
+    """
+    Align `lhs` to the Index of `rhs, if they're both pandas-like.
+
+    Notes:
+        This is only really intended for backwards-compatibility purposes,
+        for example if your library already aligns indices for users.
+        If you're designing a new library, we highly encourage you to not
+        rely on the Index.
+        For non-pandas-like inputs, this only checks that `lhs` and `rhs`
+        are the same length.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import narwhals as nw
+        >>> df_pd = pd.DataFrame({'a': [1, 2]}, index=[3, 4])
+        >>> s_pd = pd.Series([6, 7], index=[4, 3])
+        >>> df = nw.from_native(df_pd)
+        >>> s = nw.from_native(s_pd, series_only=True)
+        >>> nw.to_native(nw.maybe_align_index(df, s))
+           a
+        4  2
+        3  1
+    """
+    from narwhals._pandas_like.dataframe import PandasDataFrame
+    from narwhals._pandas_like.series import PandasSeries
+    from narwhals.dataframe import DataFrame
+    from narwhals.series import Series
+
+    def _validate_index(index: Any) -> None:
+        if not index.is_unique:
+            raise ValueError("given index doesn't have a unique index")
+
+    lhs_any = cast(Any, lhs)
+    rhs_any = cast(Any, rhs)
+    if isinstance(getattr(lhs_any, "_dataframe", None), PandasDataFrame) and isinstance(
+        getattr(rhs_any, "_dataframe", None), PandasDataFrame
+    ):
+        _validate_index(lhs_any._dataframe._dataframe.index)
+        _validate_index(rhs_any._dataframe._dataframe.index)
+        return DataFrame(  # type: ignore[return-value]
+            lhs_any._dataframe._from_dataframe(
+                lhs_any._dataframe._dataframe.loc[rhs_any._dataframe._dataframe.index]
+            )
+        )
+    if isinstance(getattr(lhs_any, "_dataframe", None), PandasDataFrame) and isinstance(
+        getattr(rhs_any, "_series", None), PandasSeries
+    ):
+        _validate_index(lhs_any._dataframe._dataframe.index)
+        _validate_index(rhs_any._series._series.index)
+        return DataFrame(  # type: ignore[return-value]
+            lhs_any._dataframe._from_dataframe(
+                lhs_any._dataframe._dataframe.loc[rhs_any._series._series.index]
+            )
+        )
+    if isinstance(getattr(lhs_any, "_series", None), PandasSeries) and isinstance(
+        getattr(rhs_any, "_dataframe", None), PandasDataFrame
+    ):
+        _validate_index(lhs_any._series._series.index)
+        _validate_index(rhs_any._dataframe._dataframe.index)
+        return Series(  # type: ignore[return-value]
+            lhs_any._series._from_series(
+                lhs_any._series._series.loc[rhs_any._dataframe._dataframe.index]
+            )
+        )
+    if isinstance(getattr(lhs_any, "_series", None), PandasSeries) and isinstance(
+        getattr(rhs_any, "_series", None), PandasSeries
+    ):
+        _validate_index(lhs_any._series._series.index)
+        _validate_index(rhs_any._series._series.index)
+        return Series(  # type: ignore[return-value]
+            lhs_any._series._from_series(
+                lhs_any._series._series.loc[rhs_any._series._series.index]
+            )
+        )
+    if len(lhs_any) != len(rhs_any):
+        msg = f"Expected `lhs` and `rhs` to have the same length, got {len(lhs_any)} and {len(rhs_any)}"
+        raise ValueError(msg)
+    return lhs
diff --git a/utils/check_api_reference.py b/utils/check_api_reference.py
index ac26591d..9eac56de 100644
--- a/utils/check_api_reference.py
+++ b/utils/check_api_reference.py
@@ -133,7 +133,7 @@ if missing := set(expr).difference(series).difference({"over"}):
 if (
     extra := set(series)
     .difference(expr)
-    .difference({"to_pandas", "to_numpy", "dtype", "name", "shape"})
+    .difference({"to_pandas", "to_numpy", "dtype", "name", "shape", "to_frame"})
 ):
     print("in series but not in expr")  # noqa: T201
     print(extra)  # noqa: T201


Test patch:
diff --git a/tests/frame/len_test.py b/tests/frame/len_test.py
new file mode 100644
index 00000000..45921ae5
--- /dev/null
+++ b/tests/frame/len_test.py
@@ -0,0 +1,18 @@
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+
+data = {
+    "a": [1.0, 2.0, None, 4.0],
+    "b": [None, 3.0, None, 5.0],
+}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+def test_drop_nulls(constructor: Any) -> None:
+    result = len(nw.from_native(constructor(data)))
+    assert result == 4
diff --git a/tests/series/to_frame_test.py b/tests/series/to_frame_test.py
new file mode 100644
index 00000000..74662409
--- /dev/null
+++ b/tests/series/to_frame_test.py
@@ -0,0 +1,16 @@
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+data = [1, 2, 3]
+
+
+@pytest.mark.parametrize("constructor", [pd.Series, pl.Series])
+def test_eq_ne(constructor: Any) -> None:
+    df = nw.from_native(constructor(data), series_only=True).alias("").to_frame()
+    compare_dicts(df, {"": [1, 2, 3]})
diff --git a/tests/utils_test.py b/tests/utils_test.py
new file mode 100644
index 00000000..bdd22796
--- /dev/null
+++ b/tests/utils_test.py
@@ -0,0 +1,47 @@
+import pandas as pd
+import polars as pl
+import pytest
+from pandas.testing import assert_frame_equal
+from pandas.testing import assert_series_equal
+
+import narwhals as nw
+
+
+def test_maybe_align_index_pandas() -> None:
+    df = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}, index=[1, 2, 0]))
+    s = nw.from_native(pd.Series([1, 2, 3], index=[2, 1, 0]), series_only=True)
+    result = nw.maybe_align_index(df, s)
+    expected = pd.DataFrame({"a": [2, 1, 3]}, index=[2, 1, 0])
+    assert_frame_equal(nw.to_native(result), expected)
+    result = nw.maybe_align_index(s, df)
+    expected = pd.Series([2, 1, 3], index=[1, 2, 0])
+    assert_series_equal(nw.to_native(result), expected)
+    result = nw.maybe_align_index(s, s.sort(descending=True))
+    expected = pd.Series([3, 2, 1], index=[0, 1, 2])
+    assert_series_equal(nw.to_native(result), expected)
+    result = nw.maybe_align_index(df, df.sort("a", descending=True))
+    expected = pd.DataFrame({"a": [3, 2, 1]}, index=[0, 2, 1])
+    assert_frame_equal(nw.to_native(result), expected)
+
+
+def test_non_unique_index() -> None:
+    df = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}, index=[1, 2, 0]))
+    s = nw.from_native(pd.Series([1, 2, 3], index=[2, 2, 0]), series_only=True)
+    with pytest.raises(ValueError, match="unique"):
+        nw.maybe_align_index(df, s)
+
+
+def test_maybe_align_index_polars() -> None:
+    df = nw.from_native(pl.DataFrame({"a": [1, 2, 3]}))
+    s = nw.from_native(pl.Series([1, 2, 3]), series_only=True)
+    result = nw.maybe_align_index(df, s)
+    assert result is df
+    with pytest.raises(ValueError, match="length"):
+        nw.maybe_align_index(df, s[1:])
+
+
+def test_native_namespace() -> None:
+    df = nw.from_native(pl.DataFrame({"a": [1, 2, 3]}))
+    assert nw.get_native_namespace(df) is pl
+    df = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}))
+    assert nw.get_native_namespace(df) is pd