Repository: narwhals-dev/narwhals
Pull request #1319 (2024-11-04 09:55:23)

feat: support passing index object directly into maybe_set_index
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #1239
- Closes #

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.
[Pandas set_index docs](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.set_index.html#pandas-dataframe-set-index)
[Pandas set_index test cases](https://github.com/pandas-dev/pandas/blob/9a015d19514597ad5d1c31f566a0a2bdb17d4cc5/pandas/tests/frame/methods/test_set_index.py#L309)
I've renamed the `column_names: str | list[str]` parameter to `keys: str | Series | list[Series | str]`. I referenced the Pandas docs where which have `keys: label or array-like or list of labels/arrays`. Since this is a backwards-compatability function, I figured it is good for it to be close to the Pandas spec, and the original param name `column` doesn't cover all cases anymore.


MarcoGorelli: thanks for your PR!

this would break backwards-compatibility, so I don't think we can just rename to `keys` unfortunately
Riik: Right, so we preserve the `column_names` for backwards comparability, add index, make both optional, and check that they are never both specified. They new signature will then be 
```
def maybe_set_index(df: T, column_names: str | list[str] | None, index: Series | list[Series] | None) -> T
```
I'll give that a try! 
Riik: > Additionally, could we handle the case for which df is actually a series? That's the use case I had in mind when I opened the related issue. I.e. if the native_object is a pandas like series, then I would like to modify its index and clearly that's not possible to do unless another Series/Index/iterable is passed.
Ah yes, I see the idea now! I think that should work with the current implementation, let me add a test for it 
Riik: @FBruzzesi I've added a test case and also support for setting the index on a Series. It didn't work out of the box since pd.Series ([docs](https://pandas.pydata.org/docs/reference/api/pandas.Series.index.html)) doesn't have a `set_index()` method, but you can set the index directly by assignment. 

Also, there is the extra case where passing in `column_names` in combination with a Series is now invalid. 

I feel like this function is getting a bit unwieldy, i.e. there a lot of conditionals inside the function. Would it be an option to make a separate function `maybe_set_index_series` instead? That would clean up the logic of this function and keep it simpler. Or maybe there is other ways to clean up this function a bit more. Let me know if you have any ideas! 
FBruzzesi: > @FBruzzesi I've added a test case and also support for setting the index on a Series. It didn't work out of the box since pd.Series ([docs](https://pandas.pydata.org/docs/reference/api/pandas.Series.index.html)) doesn't have a `set_index()` method, but you can set the index directly by assignment.

Thanks! It looks good, I tagged Marco asking if [`Series.set_axis`](https://pandas.pydata.org/docs/reference/api/pandas.Series.set_axis.html) should be prefered instead.
 
> Also, there is the extra case where passing in `column_names` in combination with a Series is now invalid.
> 
> I feel like this function is getting a bit unwieldy, i.e. there a lot of conditionals inside the function. Would it be an option to make a separate function `maybe_set_index_series` instead? That would clean up the logic of this function and keep it simpler. Or maybe there is other ways to clean up this function a bit more. Let me know if you have any ideas!

It still seems pretty reasonable in scope and size. I will take a deeper look tomorrow to see if there is a way to cut some corners :) 

FBruzzesi: @Riik I just merged with main branch and pushed a few changes to:

- parse index and column_names into keys, and avoid having different branches after
- use `set_axis` for series (which I just noticed it makes old version fail - I will get back to `series.index = index`)
Riik: > @Riik I just merged with main branch and pushed a few changes to parse index and column_names into keys, so that we avoid branching out for dataframe/series in multiple places.

Thank you! I think the function reads quite nicely now, with the `keys` internally it is easier to follow the flow. 

Anything else that still needs to happen on this PR?

Patch:
diff --git a/narwhals/utils.py b/narwhals/utils.py
index e538bee1..85ec5a00 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -274,15 +274,37 @@ def maybe_get_index(obj: T) -> Any | None:
     return None
 
 
-def maybe_set_index(df: T, column_names: str | list[str]) -> T:
+def maybe_set_index(
+    obj: T,
+    column_names: str | list[str] | None = None,
+    *,
+    index: Series | list[Series] | None = None,
+) -> T:
     """
-    Set columns `columns` to be the index of `df`, if `df` is pandas-like.
+    Set the index of a DataFrame or a Series, if it's pandas-like.
+
+    Arguments:
+        obj: object for which maybe set the index (can be either a Narwhals `DataFrame`
+            or `Series`).
+        column_names: name or list of names of the columns to set as index.
+            For dataframes, only one of `column_names` and `index` can be specified but
+            not both. If `column_names` is passed and `df` is a Series, then a
+            `ValueError` is raised.
+        index: series or list of series to set as index.
+
+    Raises:
+        ValueError: If one of the following condition happens:
+
+            - none of `column_names` and `index` are provided
+            - both `column_names` and `index` are provided
+            - `column_names` is provided and `df` is a Series
 
     Notes:
-        This is only really intended for backwards-compatibility purposes,
-        for example if your library already aligns indices for users.
+        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 is a no-op.
 
     Examples:
@@ -297,15 +319,49 @@ def maybe_set_index(df: T, column_names: str | list[str]) -> T:
         4  1
         5  2
     """
-    df_any = cast(Any, df)
-    native_frame = to_native(df_any)
-    if is_pandas_like_dataframe(native_frame):
+
+    df_any = cast(Any, obj)
+    native_obj = to_native(df_any)
+
+    if column_names is not None and index is not None:
+        msg = "Only one of `column_names` or `index` should be provided"
+        raise ValueError(msg)
+
+    if not column_names and not index:
+        msg = "Either `column_names` or `index` should be provided"
+        raise ValueError(msg)
+
+    if index is not None:
+        keys = (
+            [to_native(idx, pass_through=True) for idx in index]
+            if _is_iterable(index)
+            else to_native(index, pass_through=True)
+        )
+    else:
+        keys = column_names
+
+    if is_pandas_like_dataframe(native_obj):
         return df_any._from_compliant_dataframe(  # type: ignore[no-any-return]
-            df_any._compliant_frame._from_native_frame(
-                native_frame.set_index(column_names)
-            )
+            df_any._compliant_frame._from_native_frame(native_obj.set_index(keys))
+        )
+    elif is_pandas_like_series(native_obj):
+        if column_names:
+            msg = "Cannot set index using column names on a Series"
+            raise ValueError(msg)
+
+        if (
+            df_any._compliant_series._implementation is Implementation.PANDAS
+            and df_any._compliant_series._backend_version < (1,)
+        ):  # pragma: no cover
+            native_obj = native_obj.set_axis(keys, inplace=False)
+        else:
+            native_obj = native_obj.set_axis(keys)
+
+        return df_any._from_compliant_series(  # type: ignore[no-any-return]
+            df_any._compliant_series._from_native_series(native_obj)
         )
-    return df_any  # type: ignore[no-any-return]
+    else:
+        return df_any  # type: ignore[no-any-return]
 
 
 def maybe_reset_index(obj: T) -> T:


Test patch:
diff --git a/tests/utils_test.py b/tests/utils_test.py
index fb668b4d..dc2415c8 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import string
+from typing import TYPE_CHECKING
 
 import hypothesis.strategies as st
 import pandas as pd
@@ -15,6 +16,9 @@ import narwhals.stable.v1 as nw
 from tests.utils import PANDAS_VERSION
 from tests.utils import get_module_version_as_tuple
 
+if TYPE_CHECKING:
+    from narwhals.series import Series
+
 
 def test_maybe_align_index_pandas() -> None:
     df = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}, index=[1, 2, 0]))
@@ -58,21 +62,109 @@ def test_maybe_align_index_polars() -> None:
         nw.maybe_align_index(df, s[1:])
 
 
-def test_maybe_set_index_pandas() -> None:
-    df = nw.from_native(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[1, 2, 0]))
-    result = nw.maybe_set_index(df, "b")
-    expected = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[1, 2, 0]).set_index(
-        "b"
-    )
+@pytest.mark.parametrize(
+    "column_names",
+    ["b", ["a", "b"]],
+)
+def test_maybe_set_index_pandas_column_names(
+    column_names: str | list[str] | None,
+) -> None:
+    df = nw.from_native(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    result = nw.maybe_set_index(df, column_names)
+    expected = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}).set_index(column_names)
     assert_frame_equal(nw.to_native(result), expected)
 
 
-def test_maybe_set_index_polars() -> None:
+@pytest.mark.parametrize(
+    "column_names",
+    [
+        "b",
+        ["a", "b"],
+    ],
+)
+def test_maybe_set_index_polars_column_names(
+    column_names: str | list[str] | None,
+) -> None:
+    df = nw.from_native(pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    result = nw.maybe_set_index(df, column_names)
+    assert result is df
+
+
+@pytest.mark.parametrize(
+    "native_df_or_series",
+    [pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}), pd.Series([0, 1, 2])],
+)
+@pytest.mark.parametrize(
+    ("narwhals_index", "pandas_index"),
+    [
+        (nw.from_native(pd.Series([1, 2, 0]), series_only=True), pd.Series([1, 2, 0])),
+        (
+            [
+                nw.from_native(pd.Series([0, 1, 2]), series_only=True),
+                nw.from_native(pd.Series([1, 2, 0]), series_only=True),
+            ],
+            [
+                pd.Series([0, 1, 2]),
+                pd.Series([1, 2, 0]),
+            ],
+        ),
+    ],
+)
+def test_maybe_set_index_pandas_direct_index(
+    narwhals_index: Series | list[Series] | None,
+    pandas_index: pd.Series | list[pd.Series] | None,
+    native_df_or_series: pd.DataFrame | pd.Series,
+) -> None:
+    df = nw.from_native(native_df_or_series, allow_series=True)
+    result = nw.maybe_set_index(df, index=narwhals_index)
+    if isinstance(native_df_or_series, pd.Series):
+        native_df_or_series.index = pandas_index
+        assert_series_equal(nw.to_native(result), native_df_or_series)
+    else:
+        expected = native_df_or_series.set_index(pandas_index)
+        assert_frame_equal(nw.to_native(result), expected)
+
+
+@pytest.mark.parametrize(
+    "index",
+    [
+        nw.from_native(pd.Series([1, 2, 0]), series_only=True),
+        [
+            nw.from_native(pd.Series([0, 1, 2]), series_only=True),
+            nw.from_native(pd.Series([1, 2, 0]), series_only=True),
+        ],
+    ],
+)
+def test_maybe_set_index_polars_direct_index(
+    index: Series | list[Series] | None,
+) -> None:
     df = nw.from_native(pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
-    result = nw.maybe_set_index(df, "b")
+    result = nw.maybe_set_index(df, index=index)
     assert result is df
 
 
+def test_maybe_set_index_pandas_series_column_names() -> None:
+    df = nw.from_native(pd.Series([0, 1, 2]), allow_series=True)
+    with pytest.raises(
+        ValueError, match="Cannot set index using column names on a Series"
+    ):
+        nw.maybe_set_index(df, column_names=["a"])
+
+
+def test_maybe_set_index_pandas_either_index_or_column_names() -> None:
+    df = nw.from_native(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    column_names = ["a", "b"]
+    index = nw.from_native(pd.Series([0, 1, 2]), series_only=True)
+    with pytest.raises(
+        ValueError, match="Only one of `column_names` or `index` should be provided"
+    ):
+        nw.maybe_set_index(df, column_names=column_names, index=index)
+    with pytest.raises(
+        ValueError, match="Either `column_names` or `index` should be provided"
+    ):
+        nw.maybe_set_index(df)
+
+
 def test_maybe_get_index_pandas() -> None:
     pandas_df = pd.DataFrame({"a": [1, 2, 3]}, index=[1, 2, 0])
     result = nw.maybe_get_index(nw.from_native(pandas_df))