Python: add arr.list arr.last aliases

pola-rs · Dec 9, 2021 · ac767f1 · ac767f1
1 parent 614eee9
commit ac767f1
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 0 deletions.
diff --git a/py-polars/docs/source/reference/expression.rst b/py-polars/docs/source/reference/expression.rst
@@ -288,3 +288,6 @@ The following methods are available under the `expr.arr` attribute.
     ExprListNameSpace.sort
     ExprListNameSpace.reverse
     ExprListNameSpace.unique
+    ExprListNameSpace.get
+    ExprListNameSpace.first
+    ExprListNameSpace.last
diff --git a/py-polars/docs/source/reference/series.rst b/py-polars/docs/source/reference/series.rst
@@ -248,3 +248,6 @@ The following methods are available under the `Series.arr` attribute.
     ListNameSpace.sort
     ListNameSpace.reverse
     ListNameSpace.unique
+    ListNameSpace.get
+    ListNameSpace.first
+    ListNameSpace.last
diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -2149,6 +2149,18 @@ def get(self, index: int) -> "Expr":
         """
         return wrap_expr(self._pyexpr.lst_get(index))
 
+    def first(self) -> "Expr":
+        """
+        Get the first value of the sublists.
+        """
+        return self.get(0)
+
+    def last(self) -> "Expr":
+        """
+        Get the last value of the sublists.
+        """
+        return self.get(-1)
+
 
 class ExprStringNameSpace:
     """

diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -3402,6 +3402,18 @@ def get(self, index: int) -> "Series":
         """
         return pli.select(pli.lit(wrap_s(self._s)).arr.get(index)).to_series()
 
+    def first(self) -> "Series":
+        """
+        Get the first value of the sublists.
+        """
+        return self.get(0)
+
+    def last(self) -> "Series":
+        """
+        Get the last value of the sublists.
+        """
+        return self.get(-1)
+
 
 class DateTimeNameSpace:
     """

diff --git a/py-polars/tests/test_lists.py b/py-polars/tests/test_lists.py
@@ -7,11 +7,20 @@ def test_list_arr_get() -> None:
     out = a.arr.get(0)
     expected = pl.Series("a", [1, 4, 6])
     testing.assert_series_equal(out, expected)
+    out = a.arr.first()
+    testing.assert_series_equal(out, expected)
+    out = pl.select(pl.lit(a).arr.first()).to_series()
+    testing.assert_series_equal(out, expected)
 
     out = a.arr.get(-1)
     expected = pl.Series("a", [3, 5, 9])
     testing.assert_series_equal(out, expected)
+    out = a.arr.last()
+    testing.assert_series_equal(out, expected)
+    out = pl.select(pl.lit(a).arr.last()).to_series()
+    testing.assert_series_equal(out, expected)
 
+    a = pl.Series("a", [[1, 2, 3], [4, 5], [6, 7, 8, 9]])
     out = a.arr.get(-3)
     expected = pl.Series("a", [1, None, 7])
     testing.assert_series_equal(out, expected)