fix(python): pick a consistent order for the sort options in PyIR (#1…

…6350)
pola-rs · May 21, 2024 · 7b155c4 · 7b155c4
1 parent 5e6a433
commit 7b155c4
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 11 deletions.
diff --git a/py-polars/src/lazyframe/visitor/expr_nodes.rs b/py-polars/src/lazyframe/visitor/expr_nodes.rs
@@ -188,7 +188,8 @@ pub struct Sort {
     #[pyo3(get)]
     expr: usize,
     #[pyo3(get)]
-    options: PyObject,
+    /// maintain_order, nulls_last, descending
+    options: (bool, bool, bool),
 }
 
 #[pyclass]
@@ -216,8 +217,8 @@ pub struct SortBy {
     #[pyo3(get)]
     by: Vec<usize>,
     #[pyo3(get)]
-    /// descending, nulls_last, maintain_order
-    sort_options: (Vec<bool>, bool, bool),
+    /// maintain_order, nulls_last, descending
+    sort_options: (bool, bool, Vec<bool>),
 }
 
 #[pyclass]
@@ -481,8 +482,7 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
                 options.maintain_order,
                 options.nulls_last,
                 options.descending,
-            )
-                .to_object(py),
+            ),
         }
         .into_py(py),
         AExpr::Gather {
@@ -508,9 +508,9 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
             expr: expr.0,
             by: by.iter().map(|n| n.0).collect(),
             sort_options: (
-                sort_options.descending.clone(),
-                sort_options.nulls_last,
                 sort_options.maintain_order,
+                sort_options.nulls_last,
+                sort_options.descending.clone(),
             ),
         }
         .into_py(py),

diff --git a/py-polars/src/lazyframe/visitor/nodes.rs b/py-polars/src/lazyframe/visitor/nodes.rs
@@ -141,7 +141,7 @@ pub struct Sort {
     #[pyo3(get)]
     by_column: Vec<PyExprIR>,
     #[pyo3(get)]
-    sort_options: (Vec<bool>, bool, bool),
+    sort_options: (bool, bool, Vec<bool>),
     #[pyo3(get)]
     slice: Option<(i64, usize)>,
 }
@@ -350,9 +350,9 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
             input: input.0,
             by_column: by_column.iter().map(|e| e.into()).collect(),
             sort_options: (
-                sort_options.descending.clone(),
-                sort_options.nulls_last,
                 sort_options.maintain_order,
+                sort_options.nulls_last,
+                sort_options.descending.clone(),
             ),
             slice: *slice,
         }
@@ -445,7 +445,6 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
         .into_py(py),
         IR::Distinct { input, options } => Distinct {
             input: input.0,
-            // TODO, rest of options
             options: (
                 match options.keep_strategy {
                     UniqueKeepStrategy::First => "first",