fix: prevent panic with arg_sort_by

crates/polars-lazy/src/tests/queries.rs

@@ -1048,7 +1048,7 @@ fn test_arg_sort_multiple() -> PolarsResult<()> {
     let out = df
         .clone()
         .lazy()
-        .select([arg_sort_by([col("int"), col("flt")], &[true, false])])
+        .select([arg_sort_by([col("int"), col("flt")], &[true, false]).unwrap()])
         .collect()?;
 
     assert_eq!(
@@ -1063,7 +1063,7 @@ fn test_arg_sort_multiple() -> PolarsResult<()> {
     // check if this runs
     let _out = df
         .lazy()
-        .select([arg_sort_by([col("str"), col("flt")], &[true, false])])
+        .select([arg_sort_by([col("str"), col("flt")], &[true, false]).unwrap()])
         .collect()?;
     Ok(())
 }

crates/polars-plan/src/dsl/functions/index.rs

@@ -5,12 +5,14 @@ use super::*;
 /// until duplicates are found. Once duplicates are found, the next `Series` will
 /// be used and so on.
 #[cfg(feature = "range")]
-pub fn arg_sort_by<E: AsRef<[Expr]>>(by: E, descending: &[bool]) -> Expr {
+pub fn arg_sort_by<E: AsRef<[Expr]>>(by: E, descending: &[bool]) -> PolarsResult<Expr> {
     let e = &by.as_ref()[0];
-    let name = expr_output_name(e).unwrap();
-    int_range(lit(0 as IdxSize), len().cast(IDX_DTYPE), 1, IDX_DTYPE)
-        .sort_by(by, descending)
-        .alias(name.as_ref())
+    let name = expr_output_name(e)?;
+    Ok(
+        int_range(lit(0 as IdxSize), len().cast(IDX_DTYPE), 1, IDX_DTYPE)
+            .sort_by(by, descending)
+            .alias(name.as_ref()),
+    )
 }
 
 #[cfg(feature = "arg_where")]

py-polars/src/functions/lazy.rs

@@ -57,9 +57,10 @@ pub fn rolling_cov(
 }
 
 #[pyfunction]
-pub fn arg_sort_by(by: Vec<PyExpr>, descending: Vec<bool>) -> PyExpr {
+pub fn arg_sort_by(by: Vec<PyExpr>, descending: Vec<bool>) -> PyResult<PyExpr> {
     let by = by.into_iter().map(|e| e.inner).collect::<Vec<Expr>>();
-    dsl::arg_sort_by(by, &descending).into()
+    let e = dsl::arg_sort_by(by, &descending).map_err(PyPolarsErr::from)?;
+    Ok(e.into())
 }
 
 #[pyfunction]