fix[rust]: unset sorted flag after extend (#4986)

pola-rs · Sep 26, 2022 · 936dab7 · 936dab7
1 parent 2b98efe
commit 936dab7
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 0 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/append.rs b/polars/polars-core/src/chunked_array/ops/append.rs
@@ -53,6 +53,7 @@ impl ListChunked {
         let len = self.len();
         self.length += other.length;
         new_chunks(&mut self.chunks, &other.chunks, len);
+        self.set_sorted2(IsSorted::Not);
         Ok(())
     }
 }
@@ -62,6 +63,7 @@ impl<T: PolarsObject> ObjectChunked<T> {
     pub fn append(&mut self, other: &Self) {
         let len = self.len();
         self.length += other.length;
+        self.set_sorted2(IsSorted::Not);
         new_chunks(&mut self.chunks, &other.chunks, len);
     }
 }
diff --git a/polars/polars-core/src/chunked_array/ops/extend.rs b/polars/polars-core/src/chunked_array/ops/extend.rs
@@ -2,6 +2,7 @@ use arrow::compute::concatenate::concatenate;
 use arrow::Either;
 
 use crate::prelude::*;
+use crate::series::IsSorted;
 
 fn extend_immutable(immutable: &dyn Array, chunks: &mut Vec<ArrayRef>, other_chunks: &[ArrayRef]) {
     let out = if chunks.len() == 1 {
@@ -77,6 +78,7 @@ where
             }
         }
         self.compute_len();
+        self.set_sorted2(IsSorted::Not);
     }
 }
 
@@ -114,6 +116,7 @@ impl Utf8Chunked {
             }
         }
         self.compute_len();
+        self.set_sorted2(IsSorted::Not);
     }
 }
 
@@ -152,6 +155,7 @@ impl BooleanChunked {
             }
         }
         self.compute_len();
+        self.set_sorted2(IsSorted::Not);
     }
 }
 
@@ -160,6 +164,7 @@ impl ListChunked {
     pub fn extend(&mut self, other: &Self) -> PolarsResult<()> {
         // TODO! properly implement mutation
         // this is harder because we don't know the inner type of the list
+        self.set_sorted2(IsSorted::Not);
         self.append(other)
     }
 }

diff --git a/py-polars/tests/unit/test_sort.py b/py-polars/tests/unit/test_sort.py
@@ -254,3 +254,14 @@ def test_sorted_flag_unset_by_arithmetic_4937() -> None:
         "pmax": [3.6, 3.5],
         "pmin": [3.6, 0.0],
     }
+
+
+def test_unset_sorted_flag_after_extend() -> None:
+    df1 = pl.DataFrame({"Add": [37, 41], "Batch": [48, 49]}).sort("Add")
+    df2 = pl.DataFrame({"Add": [37], "Batch": [67]}).sort("Add")
+
+    df1.extend(df2)
+    assert not df1["Add"].flags["SORTED_ASC"]
+    df = df1.groupby("Add").agg([pl.col("Batch").min()]).sort("Add")
+    assert df["Add"].flags["SORTED_ASC"]
+    assert df.to_dict(False) == {"Add": [37, 41], "Batch": [48, 49]}