fix sorting of chunked numeric arrays (#3528)

polars/polars-core/src/chunked_array/ops/sort/mod.rs

@@ -169,6 +169,7 @@ where
     } else {
         let null_count = ca.null_count();
         let len = ca.len();
+
         let mut vals = Vec::with_capacity(ca.len());
 
         if !options.nulls_last {
@@ -177,13 +178,8 @@ where
         }
 
         ca.downcast_iter().for_each(|arr| {
-            // safety: we know the iterators len
-            let iter = unsafe {
-                arr.iter()
-                    .filter_map(|v| v.copied())
-                    .trust_my_length(len - null_count)
-            };
-            vals.extend_trusted_len(iter);
+            let iter = arr.iter().filter_map(|v| v.copied());
+            vals.extend(iter);
         });
         let mut_slice = if options.nulls_last {
             &mut vals[..len - null_count]
@@ -593,6 +589,7 @@ pub(crate) fn prepare_argsort(
 #[cfg(test)]
 mod test {
     use crate::prelude::*;
+    use crate::series::ops::NullBehavior;
 
     #[test]
     fn test_argsort() {

py-polars/tests/test_queries.py

@@ -1,3 +1,5 @@
+from datetime import datetime
+
 import numpy as np
 
 import polars as pl
@@ -141,3 +143,17 @@ def test_sorted_groupby_optimization() -> None:
 
         sorted_explicit = df.groupby("a").agg(pl.count()).sort("a", reverse=reverse)
         sorted_explicit.frame_equal(sorted_implicit)
+
+
+def test_median_on_shifted_col_3522() -> None:
+    df = pl.DataFrame(
+        {
+            "foo": [
+                datetime(2022, 5, 5, 12, 31, 34),
+                datetime(2022, 5, 5, 12, 47, 1),
+                datetime(2022, 5, 6, 8, 59, 11),
+            ]
+        }
+    )
+    diffs = df.select(pl.col("foo").diff().dt.seconds())
+    assert diffs.select(pl.col("foo").median()).to_series()[0] == 36828.5