fix: correct sorted flag of chunked gather (#14570)

pola-rs · Feb 18, 2024 · bf5310d · bf5310d
1 parent e40140f
commit bf5310d
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 19 deletions.
diff --git a/crates/polars-core/src/chunked_array/ops/gather.rs b/crates/polars-core/src/chunked_array/ops/gather.rs
@@ -9,6 +9,7 @@ use crate::chunked_array::ops::{ChunkTake, ChunkTakeUnchecked};
 use crate::chunked_array::ChunkedArray;
 use crate::datatypes::{IdxCa, PolarsDataType, StaticArray};
 use crate::prelude::*;
+use crate::series::IsSorted;
 
 const BINARY_SEARCH_LIMIT: usize = 8;
 
@@ -187,6 +188,18 @@ impl NotSpecialized for DecimalType {}
 #[cfg(feature = "object")]
 impl<T> NotSpecialized for ObjectType<T> {}
 
+pub fn _update_gather_sorted_flag(sorted_arr: IsSorted, sorted_idx: IsSorted) -> IsSorted {
+    use crate::series::IsSorted::*;
+    match (sorted_arr, sorted_idx) {
+        (_, Not) => Not,
+        (Not, _) => Not,
+        (Ascending, Ascending) => Ascending,
+        (Ascending, Descending) => Descending,
+        (Descending, Ascending) => Descending,
+        (Descending, Descending) => Ascending,
+    }
+}
+
 impl<T: PolarsDataType + NotSpecialized> ChunkTakeUnchecked<IdxCa> for ChunkedArray<T> {
     /// Gather values from ChunkedArray by index.
     unsafe fn take_unchecked(&self, indices: &IdxCa) -> Self {
@@ -233,16 +246,8 @@ impl<T: PolarsDataType + NotSpecialized> ChunkTakeUnchecked<IdxCa> for ChunkedAr
         });
 
         let mut out = ChunkedArray::from_chunk_iter_like(ca, chunks);
+        let sorted_flag = _update_gather_sorted_flag(ca.is_sorted_flag(), indices.is_sorted_flag());
 
-        use crate::series::IsSorted::*;
-        let sorted_flag = match (ca.is_sorted_flag(), indices.is_sorted_flag()) {
-            (_, Not) => Not,
-            (Not, _) => Not,
-            (Ascending, Ascending) => Ascending,
-            (Ascending, Descending) => Descending,
-            (Descending, Ascending) => Descending,
-            (Descending, Descending) => Ascending,
-        };
         out.set_sorted_flag(sorted_flag);
         out
     }
@@ -262,15 +267,8 @@ impl ChunkTakeUnchecked<IdxCa> for BinaryChunked {
 
         let mut out = ChunkedArray::from_chunks(self.name(), chunks);
 
-        use crate::series::IsSorted::*;
-        let sorted_flag = match (self.is_sorted_flag(), indices.is_sorted_flag()) {
-            (_, Not) => Not,
-            (Not, _) => Not,
-            (Ascending, Ascending) => Ascending,
-            (Ascending, Descending) => Descending,
-            (Descending, Ascending) => Descending,
-            (Descending, Descending) => Ascending,
-        };
+        let sorted_flag =
+            _update_gather_sorted_flag(self.is_sorted_flag(), indices.is_sorted_flag());
         out.set_sorted_flag(sorted_flag);
         out
     }

diff --git a/crates/polars-ops/src/chunked_array/gather/chunked.rs b/crates/polars-ops/src/chunked_array/gather/chunked.rs
@@ -1,3 +1,4 @@
+use polars_core::prelude::gather::_update_gather_sorted_flag;
 use polars_core::prelude::*;
 use polars_core::series::IsSorted;
 use polars_core::with_match_physical_numeric_polars_type;
@@ -196,7 +197,8 @@ where
             let arr = iter.collect_arr_trusted_with_dtype(arrow_dtype);
             ChunkedArray::with_chunk(self.name(), arr)
         };
-        out.set_sorted_flag(sorted);
+        let sorted_flag = _update_gather_sorted_flag(self.is_sorted_flag(), sorted);
+        out.set_sorted_flag(sorted_flag);
         out
     }
 

diff --git a/py-polars/tests/unit/operations/test_sort.py b/py-polars/tests/unit/operations/test_sort.py
@@ -789,3 +789,10 @@ def test_sort_with_null_12272() -> None:
 )
 def test_sort_series_nulls_last(input: list[Any], expected: list[Any]) -> None:
     assert pl.Series(input).sort(nulls_last=True).to_list() == expected
+
+
+def test_sorted_flag_14552() -> None:
+    a = pl.DataFrame({"a": [2, 1, 3]})
+
+    a = pl.concat([a, a], rechunk=False)
+    assert not a.join(a, on="a", how="left")["a"].flags["SORTED_ASC"]