fix(python): explicit output type in apply (#5328)

py-polars/polars/internals/expr/expr.py

@@ -3019,7 +3019,7 @@ def map(
     def apply(
         self,
         f: Callable[[pli.Series], pli.Series] | Callable[[Any], Any],
-        return_dtype: type[DataType] | None = None,
+        return_dtype: PolarsDataType | None = None,
     ) -> Expr:
         """
         Apply a custom/user-defined function (UDF) in a GroupBy or Projection context.

py-polars/polars/internals/series/series.py

@@ -3264,7 +3264,7 @@ def tanh(self) -> Series:
     def apply(
         self,
         func: Callable[[Any], Any],
-        return_dtype: type[DataType] | None = None,
+        return_dtype: PolarsDataType | None = None,
         skip_nulls: bool = True,
     ) -> Series:
         """

py-polars/src/conversion.rs

@@ -24,6 +24,7 @@ use pyo3::{PyAny, PyResult};
 use crate::dataframe::PyDataFrame;
 use crate::error::PyPolarsErr;
 use crate::lazy::dataframe::PyLazyFrame;
+use crate::object::OBJECT_NAME;
 use crate::prelude::*;
 use crate::py_modules::POLARS;
 use crate::series::PySeries;
@@ -343,7 +344,7 @@ impl FromPyObject<'_> for Wrap<DataType> {
                     "Float32" => DataType::Float32,
                     "Float64" => DataType::Float64,
                     #[cfg(feature = "object")]
-                    "Object" => DataType::Object("Object"),
+                    "Object" => DataType::Object(OBJECT_NAME),
                     "List" => DataType::List(Box::new(DataType::Boolean)),
                     "Null" => DataType::Null,
                     "Unknown" => DataType::Unknown,

py-polars/src/datatypes.rs

@@ -1,6 +1,7 @@
 use polars::prelude::*;
 use pyo3::{FromPyObject, PyAny, PyResult};
 
+use crate::object::OBJECT_NAME;
 use crate::Wrap;
 
 // Don't change the order of these!
@@ -92,7 +93,7 @@ impl From<PyDataType> for DataType {
             PyDataType::Duration(tu) => Duration(tu),
             PyDataType::Time => Time,
             #[cfg(feature = "object")]
-            PyDataType::Object => Object("object"),
+            PyDataType::Object => Object(OBJECT_NAME),
             PyDataType::Categorical => Categorical(None),
             PyDataType::Struct => Struct(vec![]),
         }

py-polars/src/lazy/apply.rs

@@ -3,9 +3,9 @@ use pyo3::prelude::*;
 use pyo3::types::PyList;
 
 use crate::lazy::dsl::PyExpr;
-use crate::prelude::PyDataType;
 use crate::py_modules::POLARS;
 use crate::series::PySeries;
+use crate::Wrap;
 
 trait ToSeries {
     fn to_series(&self, py: Python, py_polars_module: &PyObject, name: &str) -> Series;
@@ -39,13 +39,6 @@ impl ToSeries for PyObject {
     }
 }
 
-fn get_output_type(obj: &PyAny) -> Option<DataType> {
-    match obj.is_none() {
-        true => None,
-        false => Some(obj.extract::<PyDataType>().unwrap().into()),
-    }
-}
-
 pub(crate) fn call_lambda_with_series(
     py: Python,
     s: Series,
@@ -122,10 +115,10 @@ pub(crate) fn binary_lambda(lambda: &PyObject, a: Series, b: Series) -> PolarsRe
 pub fn map_single(
     pyexpr: &PyExpr,
     lambda: PyObject,
-    output_type: &PyAny,
+    output_type: Option<Wrap<DataType>>,
     agg_list: bool,
 ) -> PyExpr {
-    let output_type = get_output_type(output_type);
+    let output_type = output_type.map(|wrap| wrap.0);
 
     let output_type2 = output_type.clone();
     let function = move |s: Series| {
@@ -191,11 +184,9 @@ pub fn map_mul(
     pyexpr: &[PyExpr],
     py: Python,
     lambda: PyObject,
-    output_type: &PyAny,
+    output_type: Option<Wrap<DataType>>,
     apply_groups: bool,
 ) -> PyExpr {
-    let output_type = get_output_type(output_type);
-
     // get the pypolars module
     // do the import outside of the function to prevent import side effects in a hot loop.
     let pypolars = PyModule::import(py, "polars").unwrap().to_object(py);
@@ -217,7 +208,7 @@ pub fn map_mul(
     let exprs = pyexpr.iter().map(|pe| pe.clone().inner).collect::<Vec<_>>();
 
     let output_map = GetOutput::map_field(move |fld| match output_type {
-        Some(ref dt) => Field::new(fld.name(), dt.clone()),
+        Some(ref dt) => Field::new(fld.name(), dt.0.clone()),
         None => fld.clone(),
     });
     if apply_groups {

py-polars/src/lazy/dsl.rs

@@ -1044,7 +1044,12 @@ impl PyExpr {
             .into()
     }
 
-    pub fn map(&self, lambda: PyObject, output_type: &PyAny, agg_list: bool) -> PyExpr {
+    pub fn map(
+        &self,
+        lambda: PyObject,
+        output_type: Option<Wrap<DataType>>,
+        agg_list: bool,
+    ) -> PyExpr {
         map_single(self, lambda, output_type, agg_list)
     }
 

py-polars/src/lib.rs

@@ -440,7 +440,7 @@ pub fn map_mul(
     py: Python,
     pyexpr: Vec<PyExpr>,
     lambda: PyObject,
-    output_type: &PyAny,
+    output_type: Option<Wrap<DataType>>,
     apply_groups: bool,
 ) -> PyExpr {
     lazy::map_mul(&pyexpr, py, lambda, output_type, apply_groups)

py-polars/src/object.rs

@@ -10,6 +10,8 @@ use pyo3::prelude::*;
 use crate::prelude::ObjectValue;
 use crate::Wrap;
 
+pub(crate) const OBJECT_NAME: &str = "object";
+
 pub(crate) fn register_object_builder() {
     if !registry::is_object_builder_registered() {
         let object_builder = Box::new(|name: &str, capacity: usize| {

py-polars/src/series.rs

@@ -1001,7 +1001,7 @@ impl PySeries {
                 }
                 None => return dispatch_apply!(series, apply_lambda_unknown, py, lambda),
 
-                _ => return dispatch_apply!(series, apply_lambda, py, lambda),
+                _ => return dispatch_apply!(series, apply_lambda_unknown, py, lambda),
             };
 
             Ok(PySeries::new(out))

py-polars/tests/unit/test_apply.py

@@ -261,3 +261,16 @@ def test_apply_object_dtypes() -> None:
         "is_numeric1": [True, True, False, True, True],
         "is_numeric_infer": [True, True, False, True, True],
     }
+
+
+def test_apply_explicit_list_output_type() -> None:
+    out = pl.DataFrame({"str": ["a", "b"]}).with_columns(
+        [
+            pl.col("str").apply(
+                lambda _: pl.Series([1, 2, 3]), return_dtype=pl.List(pl.Int64)
+            )
+        ]
+    )
+
+    assert out.dtypes == [pl.List(pl.Int64)]
+    assert out.to_dict(False) == {"str": [[1, 2, 3], [1, 2, 3]]}