fix: Don't remove sort if first/last strategy is set in unique (#20481)

pola-rs · Dec 27, 2024 · 3aaf4c2 · 3aaf4c2
1 parent a35c66c
commit 3aaf4c2
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/crates/polars-plan/src/plans/optimizer/set_order.rs b/crates/polars-plan/src/plans/optimizer/set_order.rs
@@ -101,7 +101,12 @@ pub(super) fn set_order_flags(
                     options.maintain_order = false;
                     continue;
                 }
-                if !options.maintain_order {
+                if matches!(
+                    options.keep_strategy,
+                    UniqueKeepStrategy::First | UniqueKeepStrategy::Last
+                ) {
+                    maintain_order_above = true;
+                } else if !options.maintain_order {
                     maintain_order_above = false;
                 }
             },

diff --git a/py-polars/tests/unit/operations/unique/test_unique.py b/py-polars/tests/unit/operations/unique/test_unique.py
@@ -175,3 +175,28 @@ def test_categorical_updated_revmap_unique_20233() -> None:
         )
 
         assert_series_equal(s.unique(), pl.Series("a", ["D"], pl.Categorical))
+
+
+def test_unique_check_order_20480() -> None:
+    df = pl.DataFrame(
+        [
+            {
+                "key": "some_key",
+                "value": "second",
+                "number": 2,
+            },
+            {
+                "key": "some_key",
+                "value": "first",
+                "number": 1,
+            },
+        ]
+    )
+    assert (
+        df.lazy()
+        .sort("key", "number")
+        .unique(subset="key", keep="first")
+        .collect()["number"]
+        .item()
+        == 1
+    )