Skip to content

Commit

Permalink
Merge pull request #148 from narwhals-dev/maybe-set-index
Browse files Browse the repository at this point in the history
Maybe set index
  • Loading branch information
MarcoGorelli authored May 12, 2024
2 parents 5b91dc1 + 8268c9a commit 59365ee
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 0 deletions.
1 change: 1 addition & 0 deletions docs/api-reference/narwhals.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ Here are the top-level functions available in Narwhals.
- get_native_namespace
- len
- maybe_align_index
- maybe_set_index
- max
- mean
- min
Expand Down
2 changes: 2 additions & 0 deletions narwhals/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
from narwhals.translate import get_native_namespace
from narwhals.translate import to_native
from narwhals.utils import maybe_align_index
from narwhals.utils import maybe_set_index

__version__ = "0.8.11"

Expand All @@ -36,6 +37,7 @@
"to_native",
"from_native",
"maybe_align_index",
"maybe_set_index",
"get_native_namespace",
"all",
"col",
Expand Down
36 changes: 36 additions & 0 deletions narwhals/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,3 +161,39 @@ def _validate_index(index: Any) -> None:
msg = f"Expected `lhs` and `rhs` to have the same length, got {len(lhs_any)} and {len(rhs_any)}"
raise ValueError(msg)
return lhs


def maybe_set_index(df: T, column_names: str | list[str]) -> T:
"""
Set columns `columns` to be the index of `df`, if `df` is 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 is a no-op.
Examples:
>>> import pandas as pd
>>> import polars as pl
>>> import narwhals as nw
>>> df_pd = pd.DataFrame({'a': [1, 2], 'b': [4, 5]})
>>> df = nw.from_native(df_pd)
>>> nw.to_native(nw.maybe_set_index(df, 'b')) # doctest: +NORMALIZE_WHITESPACE
a
b
4 1
5 2
"""
from narwhals._pandas_like.dataframe import PandasDataFrame
from narwhals.dataframe import DataFrame

df_any = cast(Any, df)
if isinstance(getattr(df_any, "_dataframe", None), PandasDataFrame):
return DataFrame( # type: ignore[return-value]
df_any._dataframe._from_dataframe(
df_any._dataframe._dataframe.set_index(column_names)
)
)
return df
15 changes: 15 additions & 0 deletions tests/utils_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,3 +45,18 @@ def test_native_namespace() -> None:
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


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"
)
assert_frame_equal(nw.to_native(result), expected)


def test_maybe_set_index_polars() -> None:
df = nw.from_native(pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
result = nw.maybe_set_index(df, "b")
assert result is df

0 comments on commit 59365ee

Please sign in to comment.