Repository: narwhals-dev/narwhals

Issue:
[Enh]: adding nth() to dataframe
### We would like to learn about your use case. For example, if this feature is needed  to adopt Narwhals in an open source project, could you please enter the link to it below?


Like [nth() in polars](https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.nth.html), it would be a nice to have as an alternative to select()

### Please describe the purpose of the new feature or describe the problem to solve.

To provide more options to select columns

### Suggest a solution if possible.

@Cheukting is going to investigate

### If you have tried alternatives, please describe them below.

_No response_

### Additional information that may help us understand your needs.

_No response_

Discussion:
Cheukting: Seems [top level nth was added in polars since 0.20.26](https://github.com/pola-rs/polars/releases/tag/py-0.20.26), some tests won't pass using older versions (e.g. 0.20.3) is it a way to limit which versions to test?
FBruzzesi: > Seems [top level nth was added in polars since 0.20.26](https://github.com/pola-rs/polars/releases/tag/py-0.20.26), some tests won't pass using older versions (e.g. 0.20.3) is it a way to limit which versions to test?

That's something we want to address and take care of. In `PolarsNamespace` you could implement `nth` function which dispatch based on polars version:

```py
def nth(self,  *indices: int | Sequence[int]) -> PolarsExpr:
    pl = get_polars()
    if self._backend_version < (0, 20, 26):  # pragma: no cover
        # here goes the workaround
        ... 
     
    else:
        return PolarsExpr(pl.nth(*indices))
```
MarcoGorelli: > I would propose providing an error message for older versions of Polars to advise users to use col instead.

agreed 👍 that, and perhaps a little note about this in the docstring?
FBruzzesi: One tiny issue I found while stress testing these changes is the fact that `from_column_indices` return anonymous expressions and therefore `.name` namespace cannot be used with it. Here is what I mean:

```py
import narwhals as nw

import pandas as pd
import polars as pl
import pyarrow as pa

data = {"a": [1, 2], "b": [-1, 0]}

@nw.narwhalify
def func(df):
    return df.select(nw.nth(1).name.suffix("_sfx"))

func(pl.DataFrame(data))
shape: (2, 1)
┌───────┐
│ b_sfx │
│ ---   │
│ i64   │
╞═══════╡
│ -1    │
│ 0     │
└───────┘

func(pd.DataFrame(data))
> ValueError: Anonymous expressions are not supported in `.name.suffix`.
Instead of `nw.all()`, try using a named expression, such as `nw.col('a', 'b')`

func(pa.table(data))
> ValueError: Anonymous expressions are not supported in `.name.suffix`.
Instead of `nw.all()`, try using a named expression, such as `nw.col('a', 'b')`
```

That's because of 
```py
...
    root_names=None,
    output_names=None,
```

I am trying to work around that

MarcoGorelli: thanks @Cheukting !

@FBruzzesi true, but the whole "root names" / "output names" mechanism is probably ripe for a refactor anyway 😄 OK with leaving that as a potential follow-up?
FBruzzesi: > whole "root names" / "output names" mechanism is probably ripe for a refactor anyway 😄 

Oh my! Getting worried for this 😂

> OK with leaving that as a potential follow-up?

Sure let's do that, let's just add a disclaimer on limitation in the meanwhile then I am ok with merging it!


Cheukting: I also agree with @MarcoGorelli I can add a warning message (or something like that) @FBruzzesi I am happy to look at the refactoring
Cheukting: I think there's no way to express the `names` as is now. For Polars the `root_name` and `output_name` is only evaluated when calling `select`
MarcoGorelli: looks like the random gif generator didn't come up with an approval dog, so here you go 

![](https://i.fluffy.cc/b6wlBbdVkX83CHpLWxHZw74WJxHVQc9S.gif)

Patch:
diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md
index 9d7ac384..d678b873 100644
--- a/docs/api-reference/narwhals.md
+++ b/docs/api-reference/narwhals.md
@@ -28,6 +28,7 @@ Here are the top-level functions available in Narwhals.
         - min
         - narwhalify
         - new_series
+        - nth
         - sum
         - sum_horizontal
         - when
diff --git a/docs/requirements-docs.txt b/docs/requirements-docs.txt
index 7cd67b47..929f3579 100644
--- a/docs/requirements-docs.txt
+++ b/docs/requirements-docs.txt
@@ -1,6 +1,7 @@
 jinja2
 markdown-exec[ansi]
 mkdocs
+mkdocs-autorefs
 mkdocs-material
 mkdocstrings[python]
 polars>=1.0.0
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index 1313f113..a5f95cf7 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -33,6 +33,7 @@ from narwhals.expr import max
 from narwhals.expr import mean
 from narwhals.expr import mean_horizontal
 from narwhals.expr import min
+from narwhals.expr import nth
 from narwhals.expr import sum
 from narwhals.expr import sum_horizontal
 from narwhals.expr import when
@@ -80,6 +81,7 @@ __all__ = [
     "max",
     "mean",
     "mean_horizontal",
+    "nth",
     "sum",
     "sum_horizontal",
     "when",
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 1aceb576..367dc9b4 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -73,6 +73,31 @@ class ArrowExpr:
             backend_version=backend_version,
         )
 
+    @classmethod
+    def from_column_indices(
+        cls: type[Self], *column_indices: int, backend_version: tuple[int, ...]
+    ) -> Self:
+        from narwhals._arrow.series import ArrowSeries
+
+        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+            return [
+                ArrowSeries(
+                    df._native_frame[column_index],
+                    name=df._native_frame.column_names[column_index],
+                    backend_version=df._backend_version,
+                )
+                for column_index in column_indices
+            ]
+
+        return cls(
+            func,
+            depth=0,
+            function_name="nth",
+            root_names=None,
+            output_names=None,
+            backend_version=backend_version,
+        )
+
     def __narwhals_namespace__(self) -> ArrowNamespace:
         from narwhals._arrow.namespace import ArrowNamespace
 
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index ed14dd33..f4c3c9ee 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -110,6 +110,13 @@ class ArrowNamespace:
             *column_names, backend_version=self._backend_version
         )
 
+    def nth(self, *column_indices: int) -> ArrowExpr:
+        from narwhals._arrow.expr import ArrowExpr
+
+        return ArrowExpr.from_column_indices(
+            *column_indices, backend_version=self._backend_version
+        )
+
     def len(self) -> ArrowExpr:
         # coverage bug? this is definitely hit
         return ArrowExpr(  # pragma: no cover
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
index cd75e87b..e1fb5f87 100644
--- a/narwhals/_dask/expr.py
+++ b/narwhals/_dask/expr.py
@@ -71,6 +71,27 @@ class DaskExpr:
             backend_version=backend_version,
         )
 
+    @classmethod
+    def from_column_indices(
+        cls: type[Self],
+        *column_indices: int,
+        backend_version: tuple[int, ...],
+    ) -> Self:
+        def func(df: DaskLazyFrame) -> list[dask_expr.Series]:
+            return [
+                df._native_frame.iloc[:, column_index] for column_index in column_indices
+            ]
+
+        return cls(
+            func,
+            depth=0,
+            function_name="nth",
+            root_names=None,
+            output_names=None,
+            returns_scalar=False,
+            backend_version=backend_version,
+        )
+
     def _from_call(
         self,
         # First argument to `call` should be `dask_expr.Series`
diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py
index 1668ee32..ead7cf10 100644
--- a/narwhals/_dask/namespace.py
+++ b/narwhals/_dask/namespace.py
@@ -71,6 +71,12 @@ class DaskNamespace:
             backend_version=self._backend_version,
         )
 
+    def nth(self, *column_indices: int) -> DaskExpr:
+        return DaskExpr.from_column_indices(
+            *column_indices,
+            backend_version=self._backend_version,
+        )
+
     def lit(self, value: Any, dtype: dtypes.DType | None) -> DaskExpr:
         def convert_if_dtype(
             series: dask_expr.Series, dtype: DType | type[DType]
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 06be5439..51269951 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -81,6 +81,33 @@ class PandasLikeExpr:
             backend_version=backend_version,
         )
 
+    @classmethod
+    def from_column_indices(
+        cls: type[Self],
+        *column_indices: int,
+        implementation: Implementation,
+        backend_version: tuple[int, ...],
+    ) -> Self:
+        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
+            return [
+                PandasLikeSeries(
+                    df._native_frame.iloc[:, column_index],
+                    implementation=df._implementation,
+                    backend_version=df._backend_version,
+                )
+                for column_index in column_indices
+            ]
+
+        return cls(
+            func,
+            depth=0,
+            function_name="nth",
+            root_names=None,
+            output_names=None,
+            implementation=implementation,
+            backend_version=backend_version,
+        )
+
     def cast(
         self,
         dtype: Any,
diff --git a/narwhals/_pandas_like/namespace.py b/narwhals/_pandas_like/namespace.py
index 37c88954..804c0e8a 100644
--- a/narwhals/_pandas_like/namespace.py
+++ b/narwhals/_pandas_like/namespace.py
@@ -112,6 +112,13 @@ class PandasLikeNamespace:
             backend_version=self._backend_version,
         )
 
+    def nth(self, *column_indices: int) -> PandasLikeExpr:
+        return PandasLikeExpr.from_column_indices(
+            *column_indices,
+            implementation=self._implementation,
+            backend_version=self._backend_version,
+        )
+
     def all(self) -> PandasLikeExpr:
         return PandasLikeExpr(
             lambda df: [
diff --git a/narwhals/_polars/namespace.py b/narwhals/_polars/namespace.py
index 48ee8ebc..be9cf7a8 100644
--- a/narwhals/_polars/namespace.py
+++ b/narwhals/_polars/namespace.py
@@ -56,6 +56,15 @@ class PolarsNamespace:
 
         return func
 
+    def nth(self, *indices: int) -> PolarsExpr:
+        from narwhals._polars.expr import PolarsExpr
+
+        pl = get_polars()
+        if self._backend_version < (1, 0, 0):  # pragma: no cover
+            msg = "`nth` is only supported for Polars>=1.0.0. Please use `col` for columns selection instead."
+            raise AttributeError(msg)
+        return PolarsExpr(pl.nth(*indices))
+
     def len(self) -> PolarsExpr:
         from narwhals._polars.expr import PolarsExpr
 
diff --git a/narwhals/expr.py b/narwhals/expr.py
index b04a471d..8ff626d5 100644
--- a/narwhals/expr.py
+++ b/narwhals/expr.py
@@ -5,6 +5,7 @@ from typing import Any
 from typing import Callable
 from typing import Iterable
 from typing import Literal
+from typing import Sequence
 
 from narwhals.dependencies import is_numpy_array
 from narwhals.utils import flatten
@@ -3741,6 +3742,61 @@ def col(*names: str | Iterable[str]) -> Expr:
     return Expr(func)
 
 
+def nth(*indices: int | Sequence[int]) -> Expr:
+    """
+    Creates an expression that references one or more columns by their index(es).
+
+    Notes:
+        `nth` is not supported for Polars version<1.0.0. Please use [`col`](/api-reference/narwhals/#narwhals.col) instead.
+
+    Arguments:
+        indices: One or more indices representing the columns to retrieve.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import pyarrow as pa
+        >>> import narwhals as nw
+        >>> data = {"a": [1, 2], "b": [3, 4]}
+        >>> df_pl = pl.DataFrame(data)
+        >>> df_pd = pd.DataFrame(data)
+        >>> df_pa = pa.table(data)
+
+        We define a dataframe-agnostic function:
+
+        >>> @nw.narwhalify
+        ... def func(df):
+        ...     return df.select(nw.nth(0) * 2)
+
+        We can then pass either pandas or polars to `func`:
+
+        >>> func(df_pd)
+           a
+        0  2
+        1  4
+        >>> func(df_pl)  # doctest: +SKIP
+        shape: (2, 1)
+        ┌─────┐
+        │ a   │
+        │ --- │
+        │ i64 │
+        ╞═════╡
+        │ 2   │
+        │ 4   │
+        └─────┘
+        >>> func(df_pa)
+        pyarrow.Table
+        a: int64
+        ----
+        a: [[2,4]]
+    """
+
+    def func(plx: Any) -> Any:
+        return plx.nth(*flatten(indices))
+
+    return Expr(func)
+
+
 # Add underscore so it doesn't conflict with builtin `all`
 def all_() -> Expr:
     """
diff --git a/narwhals/stable/v1.py b/narwhals/stable/v1.py
index 0f34c88f..cc0a42be 100644
--- a/narwhals/stable/v1.py
+++ b/narwhals/stable/v1.py
@@ -940,6 +940,57 @@ def col(*names: str | Iterable[str]) -> Expr:
     return _stableify(nw.col(*names))
 
 
+def nth(*indices: int | Sequence[int]) -> Expr:
+    """
+    Creates an expression that references one or more columns by their index(es).
+
+    Notes:
+        `nth` is not supported for Polars version<1.0.0. Please use [`col`](/api-reference/narwhals/#narwhals.col) instead.
+
+    Arguments:
+        indices: One or more indices representing the columns to retrieve.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import pyarrow as pa
+        >>> import narwhals.stable.v1 as nw
+        >>> data = {"a": [1, 2], "b": [3, 4]}
+        >>> df_pl = pl.DataFrame(data)
+        >>> df_pd = pd.DataFrame(data)
+        >>> df_pa = pa.table(data)
+
+        We define a dataframe-agnostic function:
+
+        >>> @nw.narwhalify
+        ... def func(df):
+        ...     return df.select(nw.nth(0) * 2)
+
+        We can then pass either pandas or polars to `func`:
+
+        >>> func(df_pd)
+           a
+        0  2
+        1  4
+        >>> func(df_pl)  # doctest: +SKIP
+        shape: (2, 1)
+        ┌─────┐
+        │ a   │
+        │ --- │
+        │ i64 │
+        ╞═════╡
+        │ 2   │
+        │ 4   │
+        └─────┘
+        >>> func(df_pa)
+        pyarrow.Table
+        a: int64
+        ----
+        a: [[2,4]]
+    """
+    return _stableify(nw.nth(*indices))
+
+
 def len() -> Expr:
     """
     Return the number of rows.
@@ -1887,6 +1938,7 @@ __all__ = [
     "all_horizontal",
     "any_horizontal",
     "col",
+    "nth",
     "len",
     "lit",
     "min",


Test patch:
diff --git a/tests/expr_and_series/nth_test.py b/tests/expr_and_series/nth_test.py
new file mode 100644
index 00000000..00a8b5c9
--- /dev/null
+++ b/tests/expr_and_series/nth_test.py
@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import Any
+
+import polars as pl
+import pytest
+
+import narwhals.stable.v1 as nw
+from narwhals.utils import parse_version
+from tests.utils import Constructor
+from tests.utils import compare_dicts
+
+data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.1, 8, 9]}
+
+
+@pytest.mark.parametrize(
+    ("idx", "expected"),
+    [
+        (0, {"a": [1, 3, 2]}),
+        ([0, 1], {"a": [1, 3, 2], "b": [4, 4, 6]}),
+        ([0, 2], {"a": [1, 3, 2], "z": [7.1, 8, 9]}),
+    ],
+)
+def test_nth(
+    constructor: Constructor,
+    idx: int | list[int],
+    expected: dict[str, list[int]],
+    request: Any,
+) -> None:
+    if "polars" in str(constructor) and parse_version(pl.__version__) < parse_version(
+        "1.0.0"
+    ):
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(constructor(data))
+    result = df.select(nw.nth(idx))
+    compare_dicts(result, expected)
+
+
+@pytest.mark.skipif(
+    parse_version(pl.__version__) >= parse_version("1.0.0"),
+    reason="1.0.0",
+)
+def test_nth_not_supported() -> None:  # pragma: no cover
+    df = nw.from_native(pl.DataFrame(data))
+    with pytest.raises(
+        AttributeError, match="`nth` is only supported for Polars>=1.0.0."
+    ):
+        df.select(nw.nth(0))