Repository: narwhals-dev/narwhals
Pull request #146 (2024-05-12 10:01:33)

With row index 


Patch:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index 40321868..2f03207e 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -23,5 +23,6 @@
         - to_pandas
         - unique
         - with_columns
+        - with_row_index
       show_source: false
       show_bases: false
diff --git a/docs/api-reference/lazyframe.md b/docs/api-reference/lazyframe.md
index d1ed8022..341da426 100644
--- a/docs/api-reference/lazyframe.md
+++ b/docs/api-reference/lazyframe.md
@@ -20,6 +20,7 @@
         - sort
         - unique
         - with_columns
+        - with_row_index
       show_root_heading: false
       show_source: false
       show_bases: false
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index e97c9334..d1398ddf 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -6,6 +6,7 @@ from typing import Any
 from typing import Iterable
 from typing import Literal
 
+from narwhals._pandas_like.utils import create_native_series
 from narwhals._pandas_like.utils import evaluate_into_exprs
 from narwhals._pandas_like.utils import horizontal_concat
 from narwhals._pandas_like.utils import translate_dtype
@@ -101,6 +102,16 @@ class PandasDataFrame:
     def drop_nulls(self) -> Self:
         return self._from_dataframe(self._dataframe.dropna(axis=0))
 
+    def with_row_index(self, name: str) -> Self:
+        row_index = create_native_series(
+            range(len(self._dataframe)), implementation=self._implementation
+        ).alias(name)
+        return self._from_dataframe(
+            horizontal_concat(
+                [row_index._series, self._dataframe], implementation=self._implementation
+            )
+        )
+
     def filter(
         self,
         *predicates: IntoPandasExpr | Iterable[IntoPandasExpr],
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index ffde9023..a818c120 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -113,22 +113,27 @@ def parse_into_expr(implementation: str, into_expr: IntoPandasExpr) -> PandasExp
     if isinstance(into_expr, str):
         return plx.col(into_expr)
     if (np := get_numpy()) is not None and isinstance(into_expr, np.ndarray):
-        if implementation == "pandas":
-            pd = get_pandas()
-            series = pd.Series(into_expr, name="")
-        elif implementation == "modin":
-            mpd = get_modin()
-            series = mpd.Series(into_expr, name="")
-        elif implementation == "cudf":
-            cudf = get_cudf()
-            series = cudf.Series(into_expr, name="")
-        return plx._create_expr_from_series(
-            PandasSeries(series, implementation=implementation)
-        )
+        series = create_native_series(into_expr, implementation=implementation)
+        return plx._create_expr_from_series(series)
     msg = f"Expected IntoExpr, got {type(into_expr)}"  # pragma: no cover
     raise AssertionError(msg)
 
 
+def create_native_series(iterable: Any, implementation: str) -> PandasSeries:
+    from narwhals._pandas_like.series import PandasSeries
+
+    if implementation == "pandas":
+        pd = get_pandas()
+        series = pd.Series(iterable, name="")
+    elif implementation == "modin":
+        mpd = get_modin()
+        series = mpd.Series(iterable, name="")
+    elif implementation == "cudf":
+        cudf = get_cudf()
+        series = cudf.Series(iterable, name="")
+    return PandasSeries(series, implementation=implementation)
+
+
 def evaluate_into_expr(
     df: PandasDataFrame, into_expr: IntoPandasExpr
 ) -> list[PandasSeries]:
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index c3d08634..fd1282c6 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -13,6 +13,7 @@ from narwhals.dtypes import to_narwhals_dtype
 from narwhals.translate import get_cudf
 from narwhals.translate import get_modin
 from narwhals.translate import get_pandas
+from narwhals.utils import parse_version
 from narwhals.utils import validate_same_library
 
 if TYPE_CHECKING:
@@ -81,6 +82,17 @@ class BaseFrame:
     def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self:
         return function(self, *args, **kwargs)
 
+    def with_row_index(self, name: str = "index") -> Self:
+        if self._is_polars and parse_version(get_polars().__version__) < parse_version(
+            "0.20.4"
+        ):  # pragma: no cover
+            return self._from_dataframe(
+                self._dataframe.with_row_count(name),
+            )
+        return self._from_dataframe(
+            self._dataframe.with_row_index(name),
+        )
+
     def drop_nulls(self) -> Self:
         return self._from_dataframe(
             self._dataframe.drop_nulls(),
@@ -522,6 +534,46 @@ class DataFrame(BaseFrame):
         """
         return super().drop_nulls()
 
+    def with_row_index(self, name: str = "index") -> Self:
+        """
+        Insert column which enumerates rows.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> data = {'a': [1,2,3], 'b': [4,5,6]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.with_row_index()
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars:
+
+            >>> func(df_pd)
+               index  a  b
+            0      0  1  4
+            1      1  2  5
+            2      2  3  6
+            >>> func(df_pl)
+            shape: (3, 3)
+            ┌───────┬─────┬─────┐
+            │ index ┆ a   ┆ b   │
+            │ ---   ┆ --- ┆ --- │
+            │ u32   ┆ i64 ┆ i64 │
+            ╞═══════╪═════╪═════╡
+            │ 0     ┆ 1   ┆ 4   │
+            │ 1     ┆ 2   ┆ 5   │
+            │ 2     ┆ 3   ┆ 6   │
+            └───────┴─────┴─────┘
+        """
+        return super().with_row_index(name)
+
     @property
     def schema(self) -> dict[str, DType]:
         r"""
@@ -1511,6 +1563,46 @@ class LazyFrame(BaseFrame):
         """
         return super().drop_nulls()
 
+    def with_row_index(self, name: str = "index") -> Self:
+        """
+        Insert column which enumerates rows.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> data = {'a': [1,2,3], 'b': [4,5,6]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.LazyFrame(data)
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.with_row_index()
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars:
+
+            >>> func(df_pd)
+               index  a  b
+            0      0  1  4
+            1      1  2  5
+            2      2  3  6
+            >>> func(df_pl).collect()
+            shape: (3, 3)
+            ┌───────┬─────┬─────┐
+            │ index ┆ a   ┆ b   │
+            │ ---   ┆ --- ┆ --- │
+            │ u32   ┆ i64 ┆ i64 │
+            ╞═══════╪═════╪═════╡
+            │ 0     ┆ 1   ┆ 4   │
+            │ 1     ┆ 2   ┆ 5   │
+            │ 2     ┆ 3   ┆ 6   │
+            └───────┴─────┴─────┘
+        """
+        return super().with_row_index(name)
+
     @property
     def schema(self) -> dict[str, DType]:
         r"""
diff --git a/narwhals/series.py b/narwhals/series.py
index 6879b4f2..2c883df2 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -5,6 +5,8 @@ from typing import Any
 
 from narwhals.dtypes import to_narwhals_dtype
 from narwhals.dtypes import translate_dtype
+from narwhals.translate import get_cudf
+from narwhals.translate import get_modin
 from narwhals.translate import get_pandas
 from narwhals.translate import get_polars
 
@@ -35,7 +37,22 @@ class Series:
         if (pd := get_pandas()) is not None and isinstance(series, pd.Series):
             self._series = PandasSeries(series, implementation="pandas")
             return
-        msg = f"Expected pandas or Polars Series, got: {type(series)}"  # pragma: no cover
+        if (pd := get_modin()) is not None and isinstance(
+            series, pd.Series
+        ):  # pragma: no cover
+            self._series = PandasSeries(series, implementation="modin")
+            return
+        if (pd := get_cudf()) is not None and isinstance(
+            series, pd.Series
+        ):  # pragma: no cover
+            self._series = PandasSeries(series, implementation="cudf")
+            return
+        msg = (  # pragma: no cover
+            f"Expected pandas, Polars, modin, or cuDF Series, got: {type(series)}. "
+            "If passing something which is not already a Series, but is convertible "
+            "to one, you must specify `implementation=` "
+            "(e.g. `nw.Series([1,2,3], implementation='polars')`)"
+        )
         raise TypeError(msg)  # pragma: no cover
 
     def __array__(self, *args: Any, **kwargs: Any) -> np.ndarray:


Test patch:
diff --git a/tests/frame/with_row_index_test.py b/tests/frame/with_row_index_test.py
new file mode 100644
index 00000000..a7eb531a
--- /dev/null
+++ b/tests/frame/with_row_index_test.py
@@ -0,0 +1,23 @@
+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 = {
+    "a": ["foo", "bars"],
+    "ab": ["foo", "bars"],
+}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+def test_with_row_index(constructor: Any) -> None:
+    result = nw.from_native(constructor(data)).with_row_index()
+    expected = {"a": ["foo", "bars"], "ab": ["foo", "bars"], "index": [0, 1]}
+    compare_dicts(result, expected)
+    result = nw.from_native(constructor(data)).lazy().with_row_index("foo")
+    expected = {"a": ["foo", "bars"], "ab": ["foo", "bars"], "foo": [0, 1]}
+    compare_dicts(result, expected)