fix chunks

pola-rs · Jun 7, 2024 · 2b185c6 · 2b185c6
1 parent 0460731
commit 2b185c6
Showing 2 changed files with 18 additions and 3 deletions.
diff --git a/crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs b/crates/polars-core/src/chunked_array/ops/sort/arg_sort.rs
@@ -1,5 +1,3 @@
-use polars_utils::iter::EnumerateIdxTrait;
-
 use super::*;
 
 // Reduce monomorphisation.
@@ -83,8 +81,13 @@ where
 {
     let mut vals = Vec::with_capacity(len);
 
+    let mut count: IdxSize = 0;
     for arr_iter in iters {
-        vals.extend(arr_iter.into_iter().enumerate_idx());
+        vals.extend(arr_iter.into_iter().map(|v| {
+            let idx = count;
+            count += 1;
+            (idx, v)
+        }));
     }
 
     sort_impl(vals.as_mut_slice(), options);

diff --git a/py-polars/tests/unit/operations/test_sort.py b/py-polars/tests/unit/operations/test_sort.py
@@ -1000,3 +1000,15 @@ def test_sort_nan_1942() -> None:
     end = time.time()
 
     assert (end - start) < 1.0
+
+
+def test_sort_chunked_no_nulls() -> None:
+    df = pl.DataFrame({"values": [3.0, 2.0]})
+    df = pl.concat([df, df], rechunk=False)
+
+    assert df.with_columns(pl.col("values").arg_sort())["values"].to_list() == [
+        1,
+        3,
+        0,
+        2,
+    ]