chore[python]: Rename index argument for Expr.take (#4587)

pola-rs · Aug 27, 2022 · 4edda88 · 4edda88
1 parent a6686e6
commit 4edda88
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 21 deletions.
diff --git a/py-polars/polars/internals/expr/expr.py b/py-polars/polars/internals/expr/expr.py
@@ -13,7 +13,7 @@
 from polars.internals.expr.meta import ExprMetaNameSpace
 from polars.internals.expr.string import ExprStringNameSpace
 from polars.internals.expr.struct import ExprStructNameSpace
-from polars.utils import is_expr_sequence, is_pyexpr_sequence
+from polars.utils import deprecated_alias, is_expr_sequence, is_pyexpr_sequence
 
 try:
     from polars.polars import PyExpr
@@ -1922,13 +1922,16 @@ def sort_by(
 
         return wrap_expr(self._pyexpr.sort_by(by, reverse))
 
-    def take(self, index: list[int] | Expr | pli.Series | np.ndarray[Any, Any]) -> Expr:
+    @deprecated_alias(index="indices")
+    def take(
+        self, indices: int | list[int] | Expr | pli.Series | np.ndarray[Any, Any]
+    ) -> Expr:
         """
         Take values by index.
 
         Parameters
         ----------
-        index
+        indices
             An expression that leads to a UInt32 dtyped Series.
 
         Returns
@@ -1963,16 +1966,16 @@ def take(self, index: list[int] | Expr | pli.Series | np.ndarray[Any, Any]) -> E
         └───────┴───────┘
 
         """
-        if isinstance(index, list) or (
-            _NUMPY_AVAILABLE and isinstance(index, np.ndarray)
+        if isinstance(indices, list) or (
+            _NUMPY_AVAILABLE and isinstance(indices, np.ndarray)
         ):
-            index_lit = pli.lit(pli.Series("", index, dtype=UInt32))
+            indices_lit = pli.lit(pli.Series("", indices, dtype=UInt32))
         else:
-            index_lit = pli.expr_to_lit_or_expr(
-                index,  # type: ignore[arg-type]
+            indices_lit = pli.expr_to_lit_or_expr(
+                indices,  # type: ignore[arg-type]
                 str_to_lit=False,
             )
-        return pli.wrap_expr(self._pyexpr.take(index_lit._pyexpr))
+        return pli.wrap_expr(self._pyexpr.take(indices_lit._pyexpr))
 
     def shift(self, periods: int = 1) -> Expr:
         """

diff --git a/py-polars/polars/internals/series/series.py b/py-polars/polars/internals/series/series.py
@@ -1680,7 +1680,9 @@ def unique(self, maintain_order: bool = False) -> Series:
             return pli.select(pli.lit(self).unique(maintain_order)).to_series()
         return wrap_s(self._s.unique())
 
-    def take(self, indices: np.ndarray[Any, Any] | list[int] | pli.Expr) -> Series:
+    def take(
+        self, indices: int | list[int] | pli.Expr | Series | np.ndarray[Any, Any]
+    ) -> Series:
         """
         Take values by index.
 
@@ -1701,9 +1703,7 @@ def take(self, indices: np.ndarray[Any, Any] | list[int] | pli.Expr) -> Series:
         ]
 
         """
-        if isinstance(indices, pli.Expr):
-            return pli.select(pli.lit(self).take(indices)).to_series()
-        return wrap_s(self._s.take(indices))
+        return self.to_frame().select(pli.col(self.name).take(indices)).to_series()
 
     def null_count(self) -> int:
         """Count the null values in this Series."""

diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs
@@ -606,14 +606,6 @@ impl PySeries {
         self.series.arg_max()
     }
 
-    pub fn take(&self, indices: Wrap<Vec<IdxSize>>) -> PyResult<Self> {
-        let indices = indices.0;
-        let indices = IdxCa::from_vec("", indices);
-
-        let take = self.series.take(&indices).map_err(PyPolarsErr::from)?;
-        Ok(PySeries::new(take))
-    }
-
     pub fn take_with_series(&self, indices: &PySeries) -> PyResult<Self> {
         let idx = indices.series.idx().map_err(PyPolarsErr::from)?;
         let take = self.series.take(idx).map_err(PyPolarsErr::from)?;