Skip to main content

polars_python/series/
aggregation.rs

1use DataType::*;
2use polars::prelude::*;
3use pyo3::prelude::*;
4
5use super::PySeries;
6use crate::conversion::Wrap;
7use crate::utils::EnterPolarsExt;
8
9fn scalar_to_py(scalar: PyResult<Scalar>, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
10    Wrap(scalar?.as_any_value()).into_pyobject(py)
11}
12
13#[pymethods]
14impl PySeries {
15    fn any(&self, py: Python<'_>, ignore_nulls: bool) -> PyResult<Option<bool>> {
16        py.enter_polars(|| {
17            let s = self.series.read();
18            let s = s.bool()?;
19            PolarsResult::Ok(if ignore_nulls {
20                Some(s.any())
21            } else {
22                s.any_kleene()
23            })
24        })
25    }
26
27    fn all(&self, py: Python<'_>, ignore_nulls: bool) -> PyResult<Option<bool>> {
28        py.enter_polars(|| {
29            let s = self.series.read();
30            let s = s.bool()?;
31            PolarsResult::Ok(if ignore_nulls {
32                Some(s.all())
33            } else {
34                s.all_kleene()
35            })
36        })
37    }
38
39    fn is_empty(&self, py: Python<'_>, ignore_nulls: bool) -> PyResult<bool> {
40        py.enter_polars(|| {
41            let s = self.series.read();
42            PolarsResult::Ok(if ignore_nulls {
43                s.is_full_null()
44            } else {
45                s.is_empty()
46            })
47        })
48    }
49
50    fn arg_max(&self, py: Python) -> PyResult<Option<usize>> {
51        py.enter_polars_ok(|| self.series.read().arg_max())
52    }
53
54    fn arg_min(&self, py: Python) -> PyResult<Option<usize>> {
55        py.enter_polars_ok(|| self.series.read().arg_min())
56    }
57
58    fn min<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
59        scalar_to_py(py.enter_polars(|| self.series.read().min_reduce()), py)
60    }
61
62    fn max<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
63        scalar_to_py(py.enter_polars(|| self.series.read().max_reduce()), py)
64    }
65
66    fn mean<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
67        let s = self.series.read();
68        match s.dtype() {
69            Boolean => scalar_to_py(
70                py.enter_polars(|| s.cast(&DataType::UInt8).unwrap().mean_reduce()),
71                py,
72            ),
73            // For non-numeric output types we require mean_reduce.
74            dt if dt.is_temporal() => scalar_to_py(py.enter_polars(|| s.mean_reduce()), py),
75            _ => Ok(s.mean().into_pyobject(py)?),
76        }
77    }
78
79    fn median<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
80        let s = self.series.read();
81        match s.dtype() {
82            Boolean => scalar_to_py(
83                py.enter_polars(|| s.cast(&DataType::UInt8).unwrap().median_reduce()),
84                py,
85            ),
86            // For non-numeric output types we require median_reduce.
87            dt if dt.is_temporal() => scalar_to_py(py.enter_polars(|| s.median_reduce()), py),
88            _ => Ok(s.median().into_pyobject(py)?),
89        }
90    }
91
92    fn product<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
93        scalar_to_py(py.enter_polars(|| self.series.read().product()), py)
94    }
95
96    fn quantile<'py>(
97        &self,
98        py: Python<'py>,
99        quantile: Bound<'py, PyAny>,
100        interpolation: Wrap<QuantileMethod>,
101    ) -> PyResult<Bound<'py, PyAny>> {
102        // Accept either a single float or a list of floats
103        if let Ok(q_float) = quantile.extract::<f64>() {
104            // Single quantile: use quantile_reduce
105            scalar_to_py(
106                py.enter_polars(|| self.series.read().quantile_reduce(q_float, interpolation.0)),
107                py,
108            )
109        } else if let Ok(q_list) = quantile.extract::<Vec<f64>>() {
110            // Multiple quantiles: use quantiles_reduce
111            scalar_to_py(
112                py.enter_polars(|| {
113                    self.series
114                        .read()
115                        .quantiles_reduce(&q_list, interpolation.0)
116                }),
117                py,
118            )
119        } else {
120            Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
121                "quantile must be a float or a list of floats",
122            ))
123        }
124    }
125
126    fn std<'py>(&self, py: Python<'py>, ddof: u8) -> PyResult<Bound<'py, PyAny>> {
127        scalar_to_py(py.enter_polars(|| self.series.read().std_reduce(ddof)), py)
128    }
129
130    fn var<'py>(&self, py: Python<'py>, ddof: u8) -> PyResult<Bound<'py, PyAny>> {
131        scalar_to_py(py.enter_polars(|| self.series.read().var_reduce(ddof)), py)
132    }
133
134    fn sum<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
135        scalar_to_py(py.enter_polars(|| self.series.read().sum_reduce()), py)
136    }
137
138    fn first<'py>(&self, py: Python<'py>, ignore_nulls: bool) -> PyResult<Bound<'py, PyAny>> {
139        let result = if ignore_nulls {
140            py.enter_polars_ok(|| self.series.read().first_non_null())
141        } else {
142            py.enter_polars_ok(|| self.series.read().first())
143        };
144        scalar_to_py(result, py)
145    }
146
147    fn last<'py>(&self, py: Python<'py>, ignore_nulls: bool) -> PyResult<Bound<'py, PyAny>> {
148        let result = if ignore_nulls {
149            py.enter_polars_ok(|| self.series.read().last_non_null())
150        } else {
151            py.enter_polars_ok(|| self.series.read().last())
152        };
153        scalar_to_py(result, py)
154    }
155
156    #[cfg(feature = "approx_unique")]
157    fn approx_n_unique(&self, py: Python) -> PyResult<IdxSize> {
158        py.enter_polars(|| self.series.read().approx_n_unique())
159    }
160
161    #[cfg(feature = "bitwise")]
162    fn bitwise_and<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
163        scalar_to_py(py.enter_polars(|| self.series.read().and_reduce()), py)
164    }
165
166    #[cfg(feature = "bitwise")]
167    fn bitwise_or<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
168        scalar_to_py(py.enter_polars(|| self.series.read().or_reduce()), py)
169    }
170
171    #[cfg(feature = "bitwise")]
172    fn bitwise_xor<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
173        scalar_to_py(py.enter_polars(|| self.series.read().xor_reduce()), py)
174    }
175}