polars_python/
lazygroupby.rs

1use std::sync::Arc;
2
3use polars::lazy::frame::{LazyFrame, LazyGroupBy};
4use polars::prelude::{PlanCallback, Schema};
5use polars_utils::python_function::PythonObject;
6use pyo3::prelude::*;
7
8use crate::conversion::Wrap;
9use crate::error::PyPolarsErr;
10use crate::expr::ToExprs;
11use crate::{PyExpr, PyLazyFrame};
12
13#[pyclass(frozen)]
14#[repr(transparent)]
15pub struct PyLazyGroupBy {
16    // option because we cannot get a self by value in pyo3
17    pub lgb: Option<LazyGroupBy>,
18}
19
20#[pymethods]
21impl PyLazyGroupBy {
22    fn agg(&self, aggs: Vec<PyExpr>) -> PyLazyFrame {
23        let lgb = self.lgb.clone().unwrap();
24        let aggs = aggs.to_exprs();
25        lgb.agg(aggs).into()
26    }
27
28    fn head(&self, n: usize) -> PyLazyFrame {
29        let lgb = self.lgb.clone().unwrap();
30        lgb.head(Some(n)).into()
31    }
32
33    fn tail(&self, n: usize) -> PyLazyFrame {
34        let lgb = self.lgb.clone().unwrap();
35        lgb.tail(Some(n)).into()
36    }
37
38    #[pyo3(signature = (lambda, schema))]
39    fn map_groups(&self, lambda: PyObject, schema: Option<Wrap<Schema>>) -> PyResult<PyLazyFrame> {
40        let lgb = self.lgb.clone().unwrap();
41        let schema = match schema {
42            Some(schema) => Arc::new(schema.0),
43            None => LazyFrame::from(lgb.logical_plan.clone())
44                .collect_schema()
45                .map_err(PyPolarsErr::from)?,
46        };
47
48        let function = PythonObject(lambda);
49
50        Ok(lgb.apply(PlanCallback::new_python(function), schema).into())
51    }
52}
polars_python/lazygroupby.rs

polars_python/
lazygroupby.rs
