Feat: add list.get(index: int) (#2949)

skritsotalakis · skritsotalakis · web-flow · commit 57a5dde2005f · 2025-08-14T13:30:03.000+02:00
---------

Co-authored-by: skritsotalakis &lt;skritsotalakis@user.noreply.github.com&gt;
diff --git a/docs/api-reference/expr_list.md b/docs/api-reference/expr_list.md
@@ -5,6 +5,7 @@
     options:
       members:
         - contains
+        - get
         - len
         - unique
       show_source: false
diff --git a/docs/api-reference/series_list.md b/docs/api-reference/series_list.md
@@ -5,6 +5,7 @@
     options:
       members:
         - contains
+        - get
         - len
         - unique
       show_source: false
diff --git a/narwhals/_arrow/series_list.py b/narwhals/_arrow/series_list.py
@@ -17,4 +17,8 @@ def len(self) -> ArrowSeries:
         return self.with_native(pc.list_value_length(self.native).cast(pa.uint32()))
 
     unique = not_implemented()
+
     contains = not_implemented()
+
+    def get(self, index: int) -> ArrowSeries:
+        return self.with_native(pc.list_element(self.native, index))
diff --git a/narwhals/_compliant/any_namespace.py b/narwhals/_compliant/any_namespace.py
@@ -52,7 +52,10 @@ def offset_by(self, by: str) -> CompliantT_co: ...
 
 
 class ListNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):
+    def get(self, index: int) -> CompliantT_co: ...
+
     def len(self) -> CompliantT_co: ...
+
     def unique(self) -> CompliantT_co: ...
     def contains(self, item: NonNestedLiteral) -> CompliantT_co: ...
 
diff --git a/narwhals/_compliant/expr.py b/narwhals/_compliant/expr.py
@@ -1000,6 +1000,9 @@ def unique(self) -> EagerExprT:
     def contains(self, item: NonNestedLiteral) -> EagerExprT:
         return self.compliant._reuse_series_namespace("list", "contains", item=item)
 
+    def get(self, index: int) -> EagerExprT:
+        return self.compliant._reuse_series_namespace("list", "get", index=index)
+
 
 class CompliantExprNameNamespace(  # type: ignore[misc]
     _ExprNamespace[CompliantExprT_co],
diff --git a/narwhals/_duckdb/expr_list.py b/narwhals/_duckdb/expr_list.py
@@ -33,3 +33,8 @@ def contains(self, item: NonNestedLiteral) -> DuckDBExpr:
         return self.compliant._with_elementwise(
             lambda expr: F("list_contains", expr, lit(item))
         )
+
+    def get(self, index: int) -> DuckDBExpr:
+        return self.compliant._with_elementwise(
+            lambda expr: F("list_extract", expr, lit(index + 1))
+        )
diff --git a/narwhals/_ibis/expr_list.py b/narwhals/_ibis/expr_list.py
@@ -6,6 +6,8 @@
 from narwhals._compliant.any_namespace import ListNamespace
 
 if TYPE_CHECKING:
+    import ibis.expr.types as ir
+
     from narwhals._ibis.expr import IbisExpr
     from narwhals.typing import NonNestedLiteral
 
@@ -19,3 +21,9 @@ def unique(self) -> IbisExpr:
 
     def contains(self, item: NonNestedLiteral) -> IbisExpr:
         return self.compliant._with_callable(lambda expr: expr.contains(item))
+
+    def get(self, index: int) -> IbisExpr:
+        def _get(expr: ir.ArrayColumn) -> ir.Column:
+            return expr[index]
+
+        return self.compliant._with_callable(_get)
diff --git a/narwhals/_pandas_like/series_list.py b/narwhals/_pandas_like/series_list.py
@@ -33,4 +33,10 @@ def len(self) -> PandasLikeSeries:
         return self.with_native(result.astype(dtype)).alias(self.native.name)
 
     unique = not_implemented()
+
     contains = not_implemented()
+
+    def get(self, index: int) -> PandasLikeSeries:
+        result = self.native.list[index]
+        result.name = self.native.name
+        return self.with_native(result)
diff --git a/narwhals/_polars/utils.py b/narwhals/_polars/utils.py
@@ -332,6 +332,8 @@ class PolarsListNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
     @abc.abstractmethod
     def len(self) -> CompliantT: ...
 
+    get: Method[CompliantT]
+
     unique: Method[CompliantT]
 
 
diff --git a/narwhals/_spark_like/expr_list.py b/narwhals/_spark_like/expr_list.py
@@ -27,3 +27,9 @@ def func(expr: Column) -> Column:
             return F.array_contains(expr, F.lit(item))
 
         return self.compliant._with_elementwise(func)
+
+    def get(self, index: int) -> SparkLikeExpr:
+        def _get(expr: Column) -> Column:
+            return expr.getItem(index)
+
+        return self.compliant._with_elementwise(_get)
diff --git a/narwhals/expr_list.py b/narwhals/expr_list.py
@@ -105,3 +105,46 @@ def contains(self, item: NonNestedLiteral) -> ExprT:
         return self._expr._with_elementwise(
             lambda plx: self._expr._to_compliant_expr(plx).list.contains(item)
         )
+
+    def get(self, index: int) -> ExprT:
+        """Return the value by index in each list.
+
+        Negative indices are not accepted.
+
+        Returns:
+            A new expression.
+
+        Examples:
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> df_native = pl.DataFrame({"a": [[1, 2], [3, 4, None], [None, 5]]})
+            >>> df = nw.from_native(df_native)
+            >>> df.with_columns(a_first=nw.col("a").list.get(0))
+            ┌──────────────────────────┐
+            |    Narwhals DataFrame    |
+            |--------------------------|
+            |shape: (3, 2)             |
+            |┌──────────────┬─────────┐|
+            |│ a            ┆ a_first │|
+            |│ ---          ┆ ---     │|
+            |│ list[i64]    ┆ i64     │|
+            |╞══════════════╪═════════╡|
+            |│ [1, 2]       ┆ 1       │|
+            |│ [3, 4, null] ┆ 3       │|
+            |│ [null, 5]    ┆ null    │|
+            |└──────────────┴─────────┘|
+            └──────────────────────────┘
+        """
+        if not isinstance(index, int):
+            msg = (
+                f"Index must be of type 'int'. Got type '{type(index).__name__}' instead."
+            )
+            raise TypeError(msg)
+
+        if index < 0:
+            msg = f"Index {index} is out of bounds: should be greater than or equal to 0."
+            raise ValueError(msg)
+
+        return self._expr._with_elementwise(
+            lambda plx: self._expr._to_compliant_expr(plx).list.get(index)
+        )
diff --git a/narwhals/series_list.py b/narwhals/series_list.py
@@ -84,3 +84,39 @@ def contains(self, item: NonNestedLiteral) -> SeriesT:
         return self._narwhals_series._with_compliant(
             self._narwhals_series._compliant_series.list.contains(item)
         )
+
+    def get(self, index: int) -> SeriesT:
+        """Return the value by index in each list.
+
+        Negative indices are not accepted.
+
+        Returns:
+            A new series.
+
+        Examples:
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> s_native = pl.Series([[1, 2], [3, 4, None], [None, 5]])
+            >>> s = nw.from_native(s_native, series_only=True)
+            >>> s.list.get(1).to_native()  # doctest: +NORMALIZE_WHITESPACE
+            shape: (3,)
+            Series: '' [i64]
+            [
+                    2
+                    4
+                    5
+            ]
+        """
+        if not isinstance(index, int):
+            msg = (
+                f"Index must be of type 'int'. Got type '{type(index).__name__}' instead."
+            )
+            raise TypeError(msg)
+
+        if index < 0:
+            msg = f"Index {index} is out of bounds: should be greater than or equal to 0."
+            raise ValueError(msg)
+
+        return self._narwhals_series._with_compliant(
+            self._narwhals_series._compliant_series.list.get(index)
+        )
diff --git a/tests/expr_and_series/list/get_test.py b/tests/expr_and_series/list/get_test.py
@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import re
+from typing import Any
+
+import pytest
+
+import narwhals as nw
+from tests.utils import PANDAS_VERSION, Constructor, ConstructorEager, assert_equal_data
+
+data = {"a": [[1, 2], [None, 3], [None], None]}
+
+
+@pytest.mark.parametrize(("index", "expected"), [(0, {"a": [1, None, None, None]})])
+def test_get_expr(
+    request: pytest.FixtureRequest, constructor: Constructor, index: int, expected: Any
+) -> None:
+    if any(backend in str(constructor) for backend in ("dask", "cudf")):
+        request.applymarker(pytest.mark.xfail)
+
+    if "pandas" in str(constructor) and PANDAS_VERSION < (2, 2):
+        pytest.skip()
+
+    result = nw.from_native(constructor(data)).select(
+        nw.col("a").cast(nw.List(nw.Int32())).list.get(index)
+    )
+
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize(("index", "expected"), [(0, {"a": [1, None, None, None]})])
+def test_get_series(
+    request: pytest.FixtureRequest,
+    constructor_eager: ConstructorEager,
+    index: int,
+    expected: Any,
+) -> None:
+    if "cudf" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+
+    if "pandas" in str(constructor_eager) and PANDAS_VERSION < (2, 2):
+        pytest.skip()
+
+    if (
+        constructor_eager.__name__.startswith("pandas")
+        and "pyarrow" not in constructor_eager.__name__
+    ):
+        df = nw.from_native(constructor_eager(data), eager_only=True)
+        msg = re.escape("Series must be of PyArrow List type to support list namespace.")
+        with pytest.raises(TypeError, match=msg):
+            df["a"].list.get(index)
+        return
+
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    result = df["a"].cast(nw.List(nw.Int32())).list.get(index)
+
+    assert_equal_data({"a": result}, expected)
+
+
+def test_get_expr_negative_index(
+    request: pytest.FixtureRequest, constructor: Constructor
+) -> None:
+    data = {"a": [[1, 2], [None, 3], [None], None]}
+    index = -1
+
+    if any(backend in str(constructor) for backend in ("cudf",)):
+        request.applymarker(pytest.mark.xfail)
+    if "pandas" in str(constructor) and PANDAS_VERSION < (2, 2):
+        pytest.skip()
+
+    df = nw.from_native(constructor(data))
+    msg = re.escape(
+        f"Index {index} is out of bounds: should be greater than or equal to 0."
+    )
+    with pytest.raises(ValueError, match=msg):
+        df.select(nw.col("a").cast(nw.List(nw.Int32())).list.get(index))
+
+
+def test_get_series_negative_index(
+    request: pytest.FixtureRequest, constructor_eager: ConstructorEager
+) -> None:
+    data = {"a": [[1, 2], [None, 3], [None], None]}
+    index = -1
+
+    if "cudf" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+
+    if "pandas" in str(constructor_eager) and PANDAS_VERSION < (2, 2):
+        pytest.skip()
+
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    msg = re.escape(
+        f"Index {index} is out of bounds: should be greater than or equal to 0."
+    )
+    with pytest.raises(ValueError, match=msg):
+        df["a"].list.get(index)
+
+
+def test_get_expr_non_int_index(
+    request: pytest.FixtureRequest, constructor: Constructor
+) -> None:
+    data = {"a": [[1, 2], [None, 3], [None], None], "index": [0, 1, 0, 0]}
+    index = "index"
+
+    if any(backend in str(constructor) for backend in ("cudf",)):
+        request.applymarker(pytest.mark.xfail)
+    if "pandas" in str(constructor) and PANDAS_VERSION < (2, 2):
+        pytest.skip()
+
+    df = nw.from_native(constructor(data))
+    msg = re.escape(
+        f"Index must be of type 'int'. Got type '{type(index).__name__}' instead."
+    )
+    with pytest.raises(TypeError, match=msg):
+        df.select(nw.col("a").cast(nw.List(nw.Int32())).list.get(index))  # type: ignore[arg-type]
+
+
+def test_get_series_non_int_index(
+    request: pytest.FixtureRequest, constructor_eager: ConstructorEager
+) -> None:
+    data = {"a": [[1, 2], [None, 3], [None], None], "index": [0, 1, 0, 0]}
+    index = "index"
+
+    if "cudf" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+
+    if "pandas" in str(constructor_eager) and PANDAS_VERSION < (2, 2):
+        pytest.skip()
+
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    msg = re.escape(
+        f"Index must be of type 'int'. Got type '{type(index).__name__}' instead."
+    )
+    with pytest.raises(TypeError, match=msg):
+        df["a"].list.get(index)  # type: ignore[arg-type]

Original file line number	Diff line number	Diff line change
`@@ -33,3 +33,8 @@ def contains(self, item: NonNestedLiteral) -> DuckDBExpr:`
`33`	`33`	`return self.compliant._with_elementwise(`
`34`	`34`	`lambda expr: F("list_contains", expr, lit(item))`
`35`	`35`	`)`
	`36`	`+`
	`37`	`+ def get(self, index: int) -> DuckDBExpr:`
	`38`	`+ return self.compliant._with_elementwise(`
	`39`	`+ lambda expr: F("list_extract", expr, lit(index + 1))`
	`40`	`+ )`