fix(rust, python): raise on invalid sort_by (#9262)

pola-rs · Jun 7, 2023 · 169b03b · 169b03b
1 parent df353d6
commit 169b03b
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/sortby.rs b/polars/polars-lazy/src/physical_plan/expressions/sortby.rs
@@ -197,8 +197,12 @@ impl PhysicalExpr for SortByExpr {
                                     map_sorted_indices_to_group_slice(&sorted_idx, first)
                                 }
                             };
+                            let first = new_idx.first().unwrap_or_else(|| {
+                                invalid.store(true, Ordering::Relaxed);
+                                &0
+                            });
 
-                            (new_idx[0], new_idx)
+                            (*first, new_idx)
                         })
                         .collect()
                 });
@@ -276,8 +280,12 @@ impl PhysicalExpr for SortByExpr {
                                     map_sorted_indices_to_group_slice(&sorted_idx, first)
                                 }
                             };
+                            let first = new_idx.first().unwrap_or_else(|| {
+                                invalid.store(true, Ordering::Relaxed);
+                                &0
+                            });
 
-                            (new_idx[0], new_idx)
+                            (*first, new_idx)
                         })
                         .collect()
                 });

diff --git a/py-polars/tests/unit/test_errors.py b/py-polars/tests/unit/test_errors.py
@@ -648,3 +648,14 @@ def test_overflow_msg() -> None:
         match=r"could not append value: 2147483648 of type: i64 to the builder",
     ):
         pl.DataFrame([[2**31]], [("a", pl.Int32)], orient="row")
+
+
+def test_sort_by_err_9259() -> None:
+    df = pl.DataFrame(
+        {"a": [1, 1, 1], "b": [3, 2, 1], "c": [1, 1, 2]},
+        schema={"a": pl.Float32, "b": pl.Float32, "c": pl.Float32},
+    )
+    with pytest.raises(pl.ComputeError):
+        df.lazy().groupby("c").agg(
+            [pl.col("a").sort_by(pl.col("b").filter(pl.col("b") > 100)).sum()]
+        ).collect()