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 having(&self, predicates: Vec<PyExpr>) -> PyLazyGroupBy {
23        let mut lgb = self.lgb.clone().unwrap();
24        let predicates = predicates.to_exprs();
25        for predicate in predicates.into_iter() {
26            lgb = lgb.having(predicate);
27        }
28        PyLazyGroupBy { lgb: Some(lgb) }
29    }
30
31    fn agg(&self, aggs: Vec<PyExpr>) -> PyLazyFrame {
32        let lgb = self.lgb.clone().unwrap();
33        let aggs = aggs.to_exprs();
34        lgb.agg(aggs).into()
35    }
36
37    fn head(&self, n: usize) -> PyLazyFrame {
38        let lgb = self.lgb.clone().unwrap();
39        lgb.head(Some(n)).into()
40    }
41
42    fn tail(&self, n: usize) -> PyLazyFrame {
43        let lgb = self.lgb.clone().unwrap();
44        lgb.tail(Some(n)).into()
45    }
46
47    #[pyo3(signature = (lambda, schema))]
48    fn map_groups(&self, lambda: Py<PyAny>, schema: Option<Wrap<Schema>>) -> PyResult<PyLazyFrame> {
49        let lgb = self.lgb.clone().unwrap();
50        let schema = match schema {
51            Some(schema) => Arc::new(schema.0),
52            None => LazyFrame::from(lgb.logical_plan.clone())
53                .collect_schema()
54                .map_err(PyPolarsErr::from)?,
55        };
56
57        let function = PythonObject(lambda);
58
59        Ok(lgb.apply(PlanCallback::new_python(function), schema).into())
60    }
61}
polars_python/lazygroupby.rs

polars_python/
lazygroupby.rs
