Add examples for gather with arg_sort

py-polars/polars/expr/expr.py

@@ -2141,6 +2141,7 @@ def arg_sort(self, *, descending: bool = False, nulls_last: bool = False) -> Sel
         >>> df = pl.DataFrame(
         ...     {
         ...         "a": [20, 10, 30],
+        ...         "b": [1, 2, 3],
         ...     }
         ... )
         >>> df.select(pl.col("a").arg_sort())
@@ -2154,6 +2155,20 @@ def arg_sort(self, *, descending: bool = False, nulls_last: bool = False) -> Sel
         │ 0   │
         │ 2   │
         └─────┘
+
+        Use gather to apply the arg sort to other columns.
+
+        >>> df.select(pl.col("b").gather(pl.col("a").arg_sort()))
+        shape: (3, 1)
+        ┌─────┐
+        │ b   │
+        │ --- │
+        │ i64 │
+        ╞═════╡
+        │ 2   │
+        │ 1   │
+        │ 3   │
+        └─────┘
         """
         return self._from_pyexpr(self._pyexpr.arg_sort(descending, nulls_last))
 

py-polars/polars/functions/lazy.py

@@ -1541,11 +1541,6 @@ def arg_sort_by(
     """
     Return the row indices that would sort the column(s).
 
-    The returned expression's first element is the index of the row with the
-    lowest value of `exprs` (or highest value if `descending=True`). This row
-    would be first if the dataframe were sorted on `exprs`. The second element
-    is the index of the row that would be second if sorted, and so on.
-
     Parameters
     ----------
     exprs
@@ -1570,6 +1565,7 @@ def arg_sort_by(
     ...     {
     ...         "a": [0, 1, 1, 0],
     ...         "b": [3, 2, 3, 2],
+    ...         "c": [1, 2, 3, 4],
     ...     }
     ... )
     >>> df.select(pl.arg_sort_by("a"))
@@ -1600,6 +1596,21 @@ def arg_sort_by(
     │ 0   │
     │ 3   │
     └─────┘
+
+    Use gather to apply the arg sort to other columns.
+
+    >>> df.select(pl.col("c").gather(pl.arg_sort_by("a")))
+    shape: (4, 1)
+    ┌─────┐
+    │ c   │
+    │ --- │
+    │ i64 │
+    ╞═════╡
+    │ 1   │
+    │ 4   │
+    │ 2   │
+    │ 3   │
+    └─────┘
     """
     exprs = parse_as_list_of_expressions(exprs, *more_exprs)