allow nulls last in sort by expressions (#4030)

polars/polars-lazy/src/frame/mod.rs

@@ -291,13 +291,14 @@ impl LazyFrame {
     /// /// Sort DataFrame by 'sepal.width' column
     /// fn example(df: DataFrame) -> LazyFrame {
     ///       df.lazy()
-    ///         .sort_by_exprs(vec![col("sepal.width")], vec![false])
+    ///         .sort_by_exprs(vec![col("sepal.width")], vec![false], false)
     /// }
     /// ```
     pub fn sort_by_exprs<E: AsRef<[Expr]>, B: AsRef<[bool]>>(
         self,
         by_exprs: E,
         reverse: B,
+        nulls_last: bool,
     ) -> Self {
         let by_exprs = by_exprs.as_ref().to_vec();
         let reverse = reverse.as_ref().to_vec();
@@ -307,7 +308,7 @@ impl LazyFrame {
             let opt_state = self.get_opt_state();
             let lp = self
                 .get_plan_builder()
-                .sort(by_exprs, reverse, false)
+                .sort(by_exprs, reverse, nulls_last)
                 .build();
             Self::from_logical_plan(lp, opt_state)
         }

polars/polars-lazy/src/tests/tpch.rs

@@ -80,6 +80,7 @@ fn test_q2() -> Result<()> {
         .sort_by_exprs(
             [cols(["s_acctbal", "n_name", "s_name", "p_partkey"])],
             [true, false, false, false],
+            false,
         )
         .limit(100);
 

py-polars/polars/internals/lazy_frame.py

@@ -560,17 +560,13 @@ def sort(
         nulls_last
             Place null values last. Can only be used if sorted by a single column.
         """
-        if nulls_last and not isinstance(by, str):
-            raise ValueError(
-                "nulls_last can only be combined with 'by' argument of type str"
-            )
         if type(by) is str:
             return self._from_pyldf(self._ldf.sort(by, reverse, nulls_last))
         if type(reverse) is bool:
             reverse = [reverse]
 
         by = pli.selection_to_pyexpr_list(by)
-        return self._from_pyldf(self._ldf.sort_by_exprs(by, reverse))
+        return self._from_pyldf(self._ldf.sort_by_exprs(by, reverse, nulls_last))
 
     def collect(
         self,

py-polars/src/lazy/dataframe.rs

@@ -325,10 +325,15 @@ impl PyLazyFrame {
         .into()
     }
 
-    pub fn sort_by_exprs(&self, by_column: Vec<PyExpr>, reverse: Vec<bool>) -> PyLazyFrame {
+    pub fn sort_by_exprs(
+        &self,
+        by_column: Vec<PyExpr>,
+        reverse: Vec<bool>,
+        nulls_last: bool,
+    ) -> PyLazyFrame {
         let ldf = self.ldf.clone();
         let exprs = py_exprs_to_exprs(by_column);
-        ldf.sort_by_exprs(exprs, reverse).into()
+        ldf.sort_by_exprs(exprs, reverse, nulls_last).into()
     }
     pub fn cache(&self) -> PyLazyFrame {
         let ldf = self.ldf.clone();

py-polars/tests/test_sort.py

@@ -94,8 +94,6 @@ def test_argsort_nulls() -> None:
         None,
         None,
     ]
-    with pytest.raises(ValueError):
-        a.to_frame().sort(by=["a", "b"], nulls_last=True)
 
 
 def test_argsort_window_functions() -> None:
@@ -120,3 +118,16 @@ def test_sort_nans_3740() -> None:
         }
     )
     assert df.sort("val")["key"].to_list() == [2, 4, 1, 5, 3]
+
+
+def test_sort_by_exps_nulls_last() -> None:
+    df = pl.DataFrame(
+        {
+            "a": [1, 3, -2, None, 1],
+        }
+    ).with_row_count()
+
+    assert df.sort(pl.col("a") ** 2, nulls_last=True).to_dict(False) == {
+        "row_nr": [0, 4, 2, 1, 3],
+        "a": [1, 1, -2, 3, None],
+    }