diff --git a/src/narwhals/_utils.py b/src/narwhals/_utils.py index b4d36e605c..cf73139acd 100644 --- a/src/narwhals/_utils.py +++ b/src/narwhals/_utils.py @@ -868,11 +868,17 @@ def _validate_index(index: Any) -> None: return lhs -def maybe_get_index(obj: DataFrame[Any] | LazyFrame[Any] | Series[Any]) -> Any | None: +def maybe_get_index( + obj: DataFrame[Any] | LazyFrame[Any] | Series[Any], *, as_series: bool = False +) -> Any | None: """Get the index of a DataFrame or a Series, if it's pandas-like. Arguments: obj: Dataframe or Series. + as_series: If `True`, return the index wrapped as a Narwhals `Series` + (matching the implementation of `obj`) instead of the raw native + index. This makes the result directly usable as the `index` + argument of [`maybe_set_index`][narwhals.maybe_set_index]. Notes: This is only really intended for backwards-compatibility purposes, @@ -893,11 +899,37 @@ def maybe_get_index(obj: DataFrame[Any] | LazyFrame[Any] | Series[Any]) -> Any | >>> series = nw.from_native(series_pd, series_only=True) >>> nw.maybe_get_index(series) RangeIndex(start=0, stop=2, step=1) + >>> nw.maybe_get_index(df, as_series=True) + ┌───────────────┐ + |Narwhals Series| + |---------------| + | 0 0 | + | 1 1 | + | dtype: int64 | + └───────────────┘ """ + from narwhals.translate import _from_native_impl + obj_any = cast("Any", obj) native_obj = obj_any.to_native() if is_pandas_like_dataframe(native_obj) or is_pandas_like_series(native_obj): - return native_obj.index + native_index = native_obj.index + if not as_series: + return native_index + compliant = getattr(obj_any, "_compliant_frame", None) + if compliant is None: + compliant = obj_any._compliant_series + ns = obj_any.implementation.to_native_namespace() + native_index_series = ns.Series(native_index, name=native_index.name) + return _from_native_impl( + native_index_series, + pass_through=False, + eager_only=False, + eager_or_interchange_only=False, + series_only=True, + allow_series=None, + version=compliant._version, + ) return None diff --git a/tests/utils_test.py b/tests/utils_test.py index e18e3041ab..94bb8b7599 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -197,6 +197,41 @@ def test_maybe_get_index_pandas() -> None: assert_index_equal(result_s, pandas_series.index) +def test_maybe_get_index_pandas_as_series() -> None: + pandas_df = pd.DataFrame({"a": [1, 2, 3]}, index=[1, 2, 0]) + result = nw.maybe_get_index(nw.from_native(pandas_df), as_series=True) + assert isinstance(result, nw.Series) + assert result.implementation.is_pandas() + assert_index_equal(pd.Index(result.to_native()), pandas_df.index) + + pandas_series = pd.Series([1, 2, 3], index=[1, 2, 0]) + result_s = nw.maybe_get_index( + nw.from_native(pandas_series, series_only=True), as_series=True + ) + assert isinstance(result_s, nw.Series) + assert_index_equal(pd.Index(result_s.to_native()), pandas_series.index) + + +def test_maybe_get_index_pandas_as_series_empty() -> None: + df = nw.from_native(pd.DataFrame({"a": []})) + result = nw.maybe_get_index(df, as_series=True) + assert isinstance(result, nw.Series) + assert len(result) == 0 + series = nw.from_native(pd.Series([], dtype="float64"), series_only=True) + result_s = nw.maybe_get_index(series, as_series=True) + assert isinstance(result_s, nw.Series) + assert len(result_s) == 0 + + +def test_maybe_get_index_pandas_as_series_round_trip() -> None: + like = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}, index=[7, 8, 9])) + new_series = nw.from_native(pd.Series([4, 5, 6]), series_only=True) + result = nw.maybe_set_index( + new_series, index=nw.maybe_get_index(like, as_series=True) + ) + assert_index_equal(result.to_native().index, like.to_native().index) + + def test_maybe_get_index_polars() -> None: pytest.importorskip("polars") import polars as pl @@ -209,6 +244,16 @@ def test_maybe_get_index_polars() -> None: assert result is None +def test_maybe_get_index_polars_as_series() -> None: + pytest.importorskip("polars") + import polars as pl + + df = nw.from_native(pl.DataFrame({"a": [1, 2, 3]})) + assert nw.maybe_get_index(df, as_series=True) is None + series = nw.from_native(pl.Series([1, 2, 3]), series_only=True) + assert nw.maybe_get_index(series, as_series=True) is None + + def test_maybe_reset_index_pandas() -> None: pandas_df = nw.from_native( pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[7, 8, 9])