fix[rust]: keep sorted flag after join (#4734)

pola-rs · Sep 5, 2022 · 154e1b8 · 154e1b8
1 parent 173dcd8
commit 154e1b8
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/polars/polars-core/src/frame/hash_join/mod.rs b/polars/polars-core/src/frame/hash_join/mod.rs
@@ -318,7 +318,8 @@ impl DataFrame {
         if left_join && join_tuples.len() == self.height() {
             self.clone()
         } else {
-            self._take_unchecked_slice(join_tuples, true)
+            // left join tuples are always in ascending order
+            self._take_unchecked_slice2(join_tuples, true, IsSorted::Ascending)
         }
     }
 
@@ -865,7 +866,8 @@ impl DataFrame {
         if let Some((offset, len)) = slice {
             idx = slice_slice(idx, offset, len);
         }
-        self._take_unchecked_slice(idx, true)
+        // idx from anti-semi join should alwasy be sorted
+        self._take_unchecked_slice2(idx, true, IsSorted::Ascending)
     }
 
     #[cfg(feature = "semi_anti_join")]

diff --git a/py-polars/tests/unit/test_joins.py b/py-polars/tests/unit/test_joins.py
@@ -418,3 +418,14 @@ def test_join_inline_alias_4694() -> None:
             datetime(2021, 2, 6, 0, 0),
         ],
     }
+
+
+def test_sorted_flag_after_joins() -> None:
+    a = pl.DataFrame({"a": [1, 2, 3, 4], "b": [2, 2, 1, 4]}).sort("a")
+
+    b = pl.DataFrame({"a": [1, 2, 3, 4], "b": [2, 4, 1, 4]})
+
+    for how in ["inner", "left"]:
+        assert a.join(b, how=how, on="b")["a"].flags[  # type: ignore[arg-type]
+            "SORTED_ASC"
+        ]