polars_ops/chunked_array/list/
dispersion.rs

1use super::*;
2
3pub(super) fn median_with_nulls(ca: &ListChunked) -> Series {
4    match ca.inner_dtype() {
5        DataType::Float32 => {
6            let out: Float32Chunked = ca
7                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().median().map(|v| v as f32)))
8                .with_name(ca.name().clone());
9            out.into_series()
10        },
11        #[cfg(feature = "dtype-datetime")]
12        DataType::Date => {
13            const MS_IN_DAY: i64 = 86_400_000;
14            let out: Int64Chunked = ca
15                .apply_amortized_generic(|s| {
16                    s.and_then(|s| s.as_ref().median().map(|v| (v * (MS_IN_DAY as f64)) as i64))
17                })
18                .with_name(ca.name().clone());
19            out.into_datetime(TimeUnit::Milliseconds, None)
20                .into_series()
21        },
22        dt if dt.is_temporal() => {
23            let out: Int64Chunked = ca
24                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().median().map(|v| v as i64)))
25                .with_name(ca.name().clone());
26            out.cast(dt).unwrap()
27        },
28        _ => {
29            let out: Float64Chunked = ca
30                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().median()))
31                .with_name(ca.name().clone());
32            out.into_series()
33        },
34    }
35}
36
37pub(super) fn std_with_nulls(ca: &ListChunked, ddof: u8) -> Series {
38    match ca.inner_dtype() {
39        DataType::Float32 => {
40            let out: Float32Chunked = ca
41                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().std(ddof).map(|v| v as f32)))
42                .with_name(ca.name().clone());
43            out.into_series()
44        },
45        #[cfg(feature = "dtype-duration")]
46        DataType::Duration(tu) => {
47            let out: Int64Chunked = ca
48                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().std(ddof).map(|v| v as i64)))
49                .with_name(ca.name().clone());
50            out.into_duration(*tu).into_series()
51        },
52        _ => {
53            let out: Float64Chunked = ca
54                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().std(ddof)))
55                .with_name(ca.name().clone());
56            out.into_series()
57        },
58    }
59}
60
61pub(super) fn var_with_nulls(ca: &ListChunked, ddof: u8) -> Series {
62    match ca.inner_dtype() {
63        DataType::Float32 => {
64            let out: Float32Chunked = ca
65                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().var(ddof).map(|v| v as f32)))
66                .with_name(ca.name().clone());
67            out.into_series()
68        },
69        #[cfg(feature = "dtype-duration")]
70        DataType::Duration(TimeUnit::Milliseconds) => {
71            let out: Int64Chunked = ca
72                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().var(ddof).map(|v| v as i64)))
73                .with_name(ca.name().clone());
74            out.into_duration(TimeUnit::Milliseconds).into_series()
75        },
76        #[cfg(feature = "dtype-duration")]
77        DataType::Duration(TimeUnit::Microseconds | TimeUnit::Nanoseconds) => {
78            let out: Int64Chunked = ca
79                .cast(&DataType::List(Box::new(DataType::Duration(
80                    TimeUnit::Milliseconds,
81                ))))
82                .unwrap()
83                .list()
84                .unwrap()
85                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().var(ddof).map(|v| v as i64)))
86                .with_name(ca.name().clone());
87            out.into_duration(TimeUnit::Milliseconds).into_series()
88        },
89        _ => {
90            let out: Float64Chunked = ca
91                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().var(ddof)))
92                .with_name(ca.name().clone());
93            out.into_series()
94        },
95    }
96}