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 index f9ef04e53b36..724adbebe818 100644 --- 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 index b115ccb2e6c2..0ff9c240b6b5 100644 --- 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, + ]