polars_python/map/
lazy.rs

1use polars::prelude::*;
2use pyo3::conversion::IntoPyObjectExt;
3use pyo3::prelude::*;
4use pyo3::types::PyDict;
5
6use crate::expr::datatype::PyDataTypeExpr;
7use crate::series::PySeries;
8use crate::{PyExpr, Wrap};
9
10pub(crate) fn call_lambda_with_series(
11    py: Python<'_>,
12    s: &[Column],
13    output_dtype: Option<DataType>,
14    lambda: &PyObject,
15) -> PolarsResult<Column> {
16    // Set return_dtype in kwargs
17    let dict = PyDict::new(py);
18    let output_dtype = match output_dtype {
19        None => None,
20        Some(dt) => Some(Wrap(dt).into_pyobject(py)?),
21    };
22    dict.set_item("return_dtype", output_dtype)?;
23
24    let series_objects = s
25        .iter()
26        .map(|c| PySeries::new(c.as_materialized_series().clone()).into_py_any(py))
27        .collect::<PyResult<Vec<PyObject>>>()?;
28
29    let result = lambda.call(py, (series_objects,), Some(&dict))?;
30    Ok(result
31        .extract::<PySeries>(py)
32        .map(|s| s.series.into_inner().into_column())?)
33}
34
35pub fn map_expr(
36    pyexpr: &[PyExpr],
37    lambda: PyObject,
38    output_type: Option<PyDataTypeExpr>,
39    is_elementwise: bool,
40    returns_scalar: bool,
41) -> PyExpr {
42    let output_type = output_type.map(|v| v.inner);
43    let func =
44        python_dsl::PythonUdfExpression::new(lambda, output_type, is_elementwise, returns_scalar);
45    let exprs = pyexpr.iter().map(|pe| pe.clone().inner).collect::<Vec<_>>();
46    Expr::map_many_python(exprs, func).into()
47}
polars_python/map/lazy.rs

polars_python/map/
lazy.rs
