polars_ops/chunked_array/list/
dispersion.rs

1use arrow::temporal_conversions::MICROSECONDS_IN_DAY as US_IN_DAY;
2
3use super::*;
4
5pub(super) fn median_with_nulls(ca: &ListChunked) -> Series {
6    match ca.inner_dtype() {
7        #[cfg(feature = "dtype-f16")]
8        DataType::Float16 => {
9            let out: Float16Chunked = ca
10                .apply_amortized_generic(|s| {
11                    use num_traits::FromPrimitive;
12                    use polars_utils::float16::pf16;
13
14                    s.and_then(|s| s.as_ref().median().map(|v| pf16::from_f64(v).unwrap()))
15                })
16                .with_name(ca.name().clone());
17            out.into_series()
18        },
19        DataType::Float32 => {
20            let out: Float32Chunked = ca
21                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().median().map(|v| v as f32)))
22                .with_name(ca.name().clone());
23            out.into_series()
24        },
25        #[cfg(feature = "dtype-datetime")]
26        DataType::Date => {
27            let out: Int64Chunked = ca
28                .apply_amortized_generic(|s| {
29                    s.and_then(|s| s.as_ref().median().map(|v| (v * (US_IN_DAY as f64)) as i64))
30                })
31                .with_name(ca.name().clone());
32            out.into_datetime(TimeUnit::Microseconds, None)
33                .into_series()
34        },
35        dt if dt.is_temporal() => {
36            let out: Int64Chunked = ca
37                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().median().map(|v| v as i64)))
38                .with_name(ca.name().clone());
39            out.cast(dt).unwrap()
40        },
41        _ => {
42            let out: Float64Chunked = ca
43                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().median()))
44                .with_name(ca.name().clone());
45            out.into_series()
46        },
47    }
48}
49
50pub(super) fn std_with_nulls(ca: &ListChunked, ddof: u8) -> Series {
51    match ca.inner_dtype() {
52        #[cfg(feature = "dtype-f16")]
53        DataType::Float16 => {
54            let out: Float16Chunked = ca
55                .apply_amortized_generic(|s| {
56                    use num_traits::FromPrimitive;
57                    use polars_utils::float16::pf16;
58
59                    s.and_then(|s| s.as_ref().std(ddof).map(|v| pf16::from_f64(v).unwrap()))
60                })
61                .with_name(ca.name().clone());
62            out.into_series()
63        },
64        DataType::Float32 => {
65            let out: Float32Chunked = ca
66                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().std(ddof).map(|v| v as f32)))
67                .with_name(ca.name().clone());
68            out.into_series()
69        },
70        #[cfg(feature = "dtype-duration")]
71        DataType::Duration(tu) => {
72            let out: Int64Chunked = ca
73                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().std(ddof).map(|v| v as i64)))
74                .with_name(ca.name().clone());
75            out.into_duration(*tu).into_series()
76        },
77        _ => {
78            let out: Float64Chunked = ca
79                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().std(ddof)))
80                .with_name(ca.name().clone());
81            out.into_series()
82        },
83    }
84}
85
86pub(super) fn var_with_nulls(ca: &ListChunked, ddof: u8) -> PolarsResult<Series> {
87    match ca.inner_dtype() {
88        #[cfg(feature = "dtype-f16")]
89        DataType::Float16 => {
90            let out: Float16Chunked = ca
91                .apply_amortized_generic(|s| {
92                    use num_traits::FromPrimitive;
93                    use polars_utils::float16::pf16;
94
95                    s.and_then(|s| s.as_ref().var(ddof).map(|v| pf16::from_f64(v).unwrap()))
96                })
97                .with_name(ca.name().clone());
98            Ok(out.into_series())
99        },
100        DataType::Float32 => {
101            let out: Float32Chunked = ca
102                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().var(ddof).map(|v| v as f32)))
103                .with_name(ca.name().clone());
104            Ok(out.into_series())
105        },
106        dt if dt.is_temporal() => {
107            polars_bail!(InvalidOperation: "variance of type {dt} is not supported")
108        },
109        _ => {
110            let out: Float64Chunked = ca
111                .apply_amortized_generic(|s| s.and_then(|s| s.as_ref().var(ddof)))
112                .with_name(ca.name().clone());
113            Ok(out.into_series())
114        },
115    }
116}