feat: Support serializing Expr::RenameAlias (#22988)

Author: JakubValtar

crates/polars-plan/dsl-schema.sha256

@@ -1 +1 @@
-4dbeeec6bffb6c93698545b09282c70a2c3849a3c5e36aabd19b484684269ef2
+d1baedde1bef7b2e9be40f5aae0a12f31ccbde329d03d0fab9b967f51da715d9

crates/polars-plan/src/dsl/expr/expr_dyn_fn.rs

@@ -67,29 +67,6 @@ impl Default for SpecialEq<Arc<dyn BinaryUdfOutputField>> {
     }
 }
 
-pub trait RenameAliasFn: Send + Sync {
-    fn call(&self, name: &PlSmallStr) -> PolarsResult<PlSmallStr>;
-
-    fn try_serialize(&self, _buf: &mut Vec<u8>) -> PolarsResult<()> {
-        polars_bail!(ComputeError: "serialization not supported for this renaming function")
-    }
-}
-
-impl<F> RenameAliasFn for F
-where
-    F: Fn(&PlSmallStr) -> PolarsResult<PlSmallStr> + Send + Sync,
-{
-    fn call(&self, name: &PlSmallStr) -> PolarsResult<PlSmallStr> {
-        self(name)
-    }
-}
-
-impl Debug for dyn RenameAliasFn {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "RenameAliasFn")
-    }
-}
-
 #[derive(Clone)]
 /// Wrapper type that has special equality properties
 /// depending on the inner type specialization

crates/polars-plan/src/dsl/expr/mod.rs

@@ -8,6 +8,7 @@ use polars_compute::rolling::QuantileMethod;
 use polars_core::chunked_array::cast::CastOptions;
 use polars_core::error::feature_gated;
 use polars_core::prelude::*;
+use polars_utils::format_pl_smallstr;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 #[cfg(feature = "serde")]
@@ -182,9 +183,8 @@ pub enum Expr {
     /// `Expr::Wildcard`
     /// `Expr::Exclude`
     Selector(super::selector::Selector),
-    #[cfg_attr(any(feature = "serde", feature = "dsl-schema"), serde(skip))]
     RenameAlias {
-        function: SpecialEq<Arc<dyn RenameAliasFn>>,
+        function: RenameAliasFn,
         expr: Arc<Expr>,
     },
 }
@@ -595,3 +595,51 @@ impl Operator {
         !(self.is_comparison_or_bitwise())
     }
 }
+
+#[derive(Clone, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
+pub enum RenameAliasFn {
+    Prefix(PlSmallStr),
+    Suffix(PlSmallStr),
+    ToLowercase,
+    ToUppercase,
+    #[cfg(feature = "python")]
+    Python(Arc<polars_utils::python_function::PythonObject>),
+    #[cfg_attr(any(feature = "serde", feature = "dsl-schema"), serde(skip))]
+    Rust(SpecialEq<Arc<RenameAliasRustFn>>),
+}
+
+impl RenameAliasFn {
+    pub fn call(&self, name: &PlSmallStr) -> PolarsResult<PlSmallStr> {
+        let out = match self {
+            Self::Prefix(prefix) => format_pl_smallstr!("{prefix}{name}"),
+            Self::Suffix(suffix) => format_pl_smallstr!("{name}{suffix}"),
+            Self::ToLowercase => PlSmallStr::from_string(name.to_lowercase()),
+            Self::ToUppercase => PlSmallStr::from_string(name.to_uppercase()),
+            #[cfg(feature = "python")]
+            Self::Python(lambda) => {
+                let name = name.as_str();
+                let out = pyo3::marker::Python::with_gil(|py| {
+                    let out: PlSmallStr = lambda
+                        .call1(py, (name,))?
+                        .extract::<std::borrow::Cow<str>>(py)?
+                        .as_ref()
+                        .into();
+                    pyo3::PyResult::<_>::Ok(out)
+                });
+                match out {
+                    Ok(out) => format_pl_smallstr!("{}", out),
+                    Err(e) => {
+                        polars_bail!(ComputeError: "Python function in 'name.map' produced an error: {e}.")
+                    },
+                }
+            },
+            Self::Rust(f) => f(name)?,
+        };
+        Ok(out)
+    }
+}
+
+pub type RenameAliasRustFn =
+    dyn Fn(&PlSmallStr) -> PolarsResult<PlSmallStr> + 'static + Send + Sync;

crates/polars-plan/src/dsl/name.rs

@@ -1,4 +1,3 @@
-use polars_utils::format_pl_smallstr;
 #[cfg(feature = "dtype-struct")]
 use polars_utils::pl_str::PlSmallStr;
 
@@ -30,35 +29,54 @@ impl ExprNameNameSpace {
     where
         F: Fn(&PlSmallStr) -> PolarsResult<PlSmallStr> + 'static + Send + Sync,
     {
-        let function = SpecialEq::new(Arc::new(function) as Arc<dyn RenameAliasFn>);
+        let function = SpecialEq::new(Arc::new(function) as Arc<RenameAliasRustFn>);
         Expr::RenameAlias {
             expr: Arc::new(self.0),
-            function,
+            function: RenameAliasFn::Rust(function),
+        }
+    }
+
+    /// Define an alias by mapping a python lambda over the original root column name.
+    #[cfg(feature = "python")]
+    pub fn map_udf(self, function: polars_utils::python_function::PythonObject) -> Expr {
+        Expr::RenameAlias {
+            expr: Arc::new(self.0),
+            function: RenameAliasFn::Python(Arc::new(function)),
         }
     }
 
     /// Add a prefix to the root column name.
     pub fn prefix(self, prefix: &str) -> Expr {
-        let prefix = prefix.to_string();
-        self.map(move |name| Ok(format_pl_smallstr!("{prefix}{name}")))
+        Expr::RenameAlias {
+            expr: Arc::new(self.0),
+            function: RenameAliasFn::Prefix(prefix.into()),
+        }
     }
 
     /// Add a suffix to the root column name.
     pub fn suffix(self, suffix: &str) -> Expr {
-        let suffix = suffix.to_string();
-        self.map(move |name| Ok(format_pl_smallstr!("{name}{suffix}")))
+        Expr::RenameAlias {
+            expr: Arc::new(self.0),
+            function: RenameAliasFn::Suffix(suffix.into()),
+        }
     }
 
     /// Update the root column name to use lowercase characters.
     #[allow(clippy::wrong_self_convention)]
     pub fn to_lowercase(self) -> Expr {
-        self.map(move |name| Ok(PlSmallStr::from_string(name.to_lowercase())))
+        Expr::RenameAlias {
+            expr: Arc::new(self.0),
+            function: RenameAliasFn::ToLowercase,
+        }
     }
 
     /// Update the root column name to use uppercase characters.
     #[allow(clippy::wrong_self_convention)]
     pub fn to_uppercase(self) -> Expr {
-        self.map(move |name| Ok(PlSmallStr::from_string(name.to_uppercase())))
+        Expr::RenameAlias {
+            expr: Arc::new(self.0),
+            function: RenameAliasFn::ToUppercase,
+        }
     }
 
     #[cfg(feature = "dtype-struct")]

crates/polars-plan/src/dsl/plan.rs

@@ -48,7 +48,7 @@ use super::*;
 // - changing a name, type, or meaning of a field or an enum variant
 // - changing a default value of a field or a default enum variant
 // - restricting the range of allowed values a field can have
-pub static DSL_VERSION: (u16, u16) = (8, 0);
+pub static DSL_VERSION: (u16, u16) = (8, 1);
 static DSL_MAGIC_BYTES: &[u8] = b"DSL_VERSION";
 
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

crates/polars-python/src/expr/name.rs

@@ -1,8 +1,7 @@
 use std::borrow::Cow;
 
 use polars::prelude::*;
-use polars_utils::format_pl_smallstr;
-use polars_utils::pl_str::PlSmallStr;
+use polars_utils::python_function::PythonObject;
 use pyo3::prelude::*;
 
 use crate::PyExpr;
@@ -17,15 +16,7 @@ impl PyExpr {
         self.inner
             .clone()
             .name()
-            .map(move |name| {
-                let out = Python::with_gil(|py| lambda.call1(py, (name.as_str(),)));
-                match out {
-                    Ok(out) => Ok(format_pl_smallstr!("{}", out)),
-                    Err(e) => Err(PolarsError::ComputeError(
-                        format!("Python function in 'name.map' produced an error: {e}.").into(),
-                    )),
-                }
-            })
+            .map_udf(PythonObject(lambda))
             .into()
     }
 

py-polars/tests/unit/cloud/test_prepare_cloud_plan.py

@@ -93,16 +93,11 @@ def test_prepare_cloud_plan_fail_on_local_data_source(lf: pl.LazyFrame) -> None:
 @pytest.mark.parametrize(
     "lf",
     [
-        pl.LazyFrame({"a": [{"x": 1, "y": 2}]})
-        .select(pl.col("a").name.map(lambda x: x.upper()))
-        .sink_parquet(DST, lazy=True),
         pl.LazyFrame({"a": [{"x": 1, "y": 2}]})
         .select(pl.col("a").name.map_fields(lambda x: x.upper()))
         .sink_parquet(DST, lazy=True),
     ],
 )
 def test_prepare_cloud_plan_fail_on_serialization(lf: pl.LazyFrame) -> None:
-    with pytest.raises(
-        ComputeError, match="cannot be serialized|serialization not supported"
-    ):
+    with pytest.raises(ComputeError, match="serialization not supported"):
         prepare_cloud_plan(lf)