fix: Always use cloudpickle for the python objects in cloud plans (#23474)

Author: coastalwhite

crates/polars-plan/src/client/mod.rs

@@ -2,7 +2,7 @@ mod check;
 
 use polars_core::error::PolarsResult;
 
-use crate::dsl::DslPlan;
+use crate::dsl::{DslPlan, PlanSerializationContext};
 
 /// Prepare the given [`DslPlan`] for execution on Polars Cloud.
 pub fn prepare_cloud_plan(dsl: DslPlan) -> PolarsResult<Vec<u8>> {
@@ -11,7 +11,12 @@ pub fn prepare_cloud_plan(dsl: DslPlan) -> PolarsResult<Vec<u8>> {
 
     // Serialize the plan.
     let mut writer = Vec::new();
-    dsl.serialize_versioned(&mut writer)?;
+    dsl.serialize_versioned(
+        &mut writer,
+        PlanSerializationContext {
+            use_cloudpickle: true,
+        },
+    )?;
 
     Ok(writer)
 }

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

@@ -234,6 +234,11 @@ impl Default for DslPlan {
     }
 }
 
+#[derive(Default, Clone, Copy)]
+pub struct PlanSerializationContext {
+    pub use_cloudpickle: bool,
+}
+
 impl DslPlan {
     pub fn describe(&self) -> PolarsResult<String> {
         Ok(self.clone().to_alp()?.describe())
@@ -269,9 +274,19 @@ impl DslPlan {
     }
 
     #[cfg(feature = "serde")]
-    pub fn serialize_versioned<W: Write>(&self, mut writer: W) -> PolarsResult<()> {
+    pub fn serialize_versioned<W: Write>(
+        &self,
+        mut writer: W,
+        ctx: PlanSerializationContext,
+    ) -> PolarsResult<()> {
         let le_major = DSL_VERSION.0.to_le_bytes();
         let le_minor = DSL_VERSION.1.to_le_bytes();
+
+        // @GB:
+        // This is absolute horrendous but serde does not allow for state to passed along with the
+        // serialization so there is no proper way to do this except replace serde.
+        polars_utils::pl_serialize::USE_CLOUDPICKLE.set(ctx.use_cloudpickle);
+
         writer.write_all(DSL_MAGIC_BYTES)?;
         writer.write_all(&le_major)?;
         writer.write_all(&le_minor)?;

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

@@ -9,8 +9,6 @@ use polars_core::prelude::UnknownKind;
 use polars_core::schema::Schema;
 use polars_utils::pl_str::PlSmallStr;
 use pyo3::prelude::*;
-use pyo3::pybacked::PyBackedBytes;
-use pyo3::types::PyBytes;
 
 use crate::prelude::*;
 
@@ -54,49 +52,27 @@ impl PythonUdfExpression {
 
     #[cfg(feature = "serde")]
     pub(crate) fn try_deserialize(buf: &[u8]) -> PolarsResult<Arc<dyn ColumnsUdf>> {
-        // Handle byte mark
-
         use polars_utils::pl_serialize;
-        debug_assert!(buf.starts_with(PYTHON_SERDE_MAGIC_BYTE_MARK));
-        let buf = &buf[PYTHON_SERDE_MAGIC_BYTE_MARK.len()..];
 
-        // Handle pickle metadata
-        let use_cloudpickle = buf[0];
-        if use_cloudpickle != 0 {
-            let ser_py_version = &buf[1..3];
-            let cur_py_version = *PYTHON3_VERSION;
-            polars_ensure!(
-                ser_py_version == cur_py_version,
-                InvalidOperation:
-                "current Python version {:?} does not match the Python version used to serialize the UDF {:?}",
-                (3, cur_py_version[0], cur_py_version[1]),
-                (3, ser_py_version[0], ser_py_version[1] )
-            );
+        if !buf.starts_with(PYTHON_SERDE_MAGIC_BYTE_MARK) {
+            polars_bail!(InvalidOperation: "serialization expected python magic byte mark");
         }
-        let buf = &buf[3..];
+        let buf = &buf[PYTHON_SERDE_MAGIC_BYTE_MARK.len()..];
 
         // Load UDF metadata
         let mut reader = Cursor::new(buf);
         let (output_type, is_elementwise, returns_scalar): (Option<DataTypeExpr>, bool, bool) =
             pl_serialize::deserialize_from_reader::<_, _, true>(&mut reader)?;
 
-        let remainder = &buf[reader.position() as usize..];
-
-        // Load UDF
-        Python::with_gil(|py| {
-            let pickle = PyModule::import(py, "pickle")
-                .expect("unable to import 'pickle'")
-                .getattr("loads")
-                .unwrap();
-            let arg = (PyBytes::new(py, remainder),);
-            let python_function = pickle.call1(arg)?;
-            Ok(Arc::new(Self::new(
-                python_function.into(),
-                output_type,
-                is_elementwise,
-                returns_scalar,
-            )) as Arc<dyn ColumnsUdf>)
-        })
+        let buf = &buf[reader.position() as usize..];
+        let python_function = pl_serialize::python_object_deserialize(buf)?;
+
+        Ok(Arc::new(Self::new(
+            python_function,
+            output_type,
+            is_elementwise,
+            returns_scalar,
+        )))
     }
 }
 
@@ -145,48 +121,23 @@ impl ColumnsUdf for PythonUdfExpression {
 
     #[cfg(feature = "serde")]
     fn try_serialize(&self, buf: &mut Vec<u8>) -> PolarsResult<()> {
-        // Write byte marks
-
         use polars_utils::pl_serialize;
+
+        // Write byte marks
         buf.extend_from_slice(PYTHON_SERDE_MAGIC_BYTE_MARK);
 
-        Python::with_gil(|py| {
-            // Try pickle to serialize the UDF, otherwise fall back to cloudpickle.
-            let pickle = PyModule::import(py, "pickle")
-                .expect("unable to import 'pickle'")
-                .getattr("dumps")
-                .unwrap();
-            let pickle_result = pickle.call1((self.python_function.clone_ref(py),));
-            let (dumped, use_cloudpickle) = match pickle_result {
-                Ok(dumped) => (dumped, false),
-                Err(_) => {
-                    let cloudpickle = PyModule::import(py, "cloudpickle")?
-                        .getattr("dumps")
-                        .unwrap();
-                    let dumped = cloudpickle.call1((self.python_function.clone_ref(py),))?;
-                    (dumped, true)
-                },
-            };
-
-            // Write pickle metadata
-            buf.push(use_cloudpickle as u8);
-            buf.extend_from_slice(&*PYTHON3_VERSION);
-
-            // Write UDF metadata
-            pl_serialize::serialize_into_writer::<_, _, true>(
-                &mut *buf,
-                &(
-                    self.output_type.clone(),
-                    self.is_elementwise,
-                    self.returns_scalar,
-                ),
-            )?;
-
-            // Write UDF
-            let dumped = dumped.extract::<PyBackedBytes>().unwrap();
-            buf.extend_from_slice(&dumped);
-            Ok(())
-        })
+        // Write UDF metadata
+        pl_serialize::serialize_into_writer::<_, _, true>(
+            &mut *buf,
+            &(
+                self.output_type.clone(),
+                self.is_elementwise,
+                self.returns_scalar,
+            ),
+        )?;
+
+        pl_serialize::python_object_serialize(&self.python_function, buf)?;
+        Ok(())
     }
 }
 

crates/polars-python/src/lazyframe/serde.rs

@@ -15,7 +15,11 @@ impl PyLazyFrame {
     fn serialize_binary(&self, py: Python<'_>, py_f: PyObject) -> PyResult<()> {
         let file = get_file_like(py_f, true)?;
         let writer = BufWriter::new(file);
-        py.enter_polars(|| self.ldf.logical_plan.serialize_versioned(writer))
+        py.enter_polars(|| {
+            self.ldf
+                .logical_plan
+                .serialize_versioned(writer, Default::default())
+        })
     }
 
     /// Serialize into a JSON string.

crates/polars-utils/src/pl_serialize.rs

@@ -215,6 +215,81 @@ where
     }
 }
 
+thread_local! {
+    pub static USE_CLOUDPICKLE: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
+}
+
+#[cfg(feature = "python")]
+pub fn python_object_serialize(
+    pyobj: &pyo3::Py<pyo3::PyAny>,
+    buf: &mut Vec<u8>,
+) -> PolarsResult<()> {
+    use pyo3::Python;
+    use pyo3::pybacked::PyBackedBytes;
+    use pyo3::types::{PyAnyMethods, PyModule};
+
+    use crate::python_function::PYTHON3_VERSION;
+
+    let mut use_cloudpickle = USE_CLOUDPICKLE.get();
+    let dumped = Python::with_gil(|py| {
+        // Pickle with whatever pickling method was selected.
+        if use_cloudpickle {
+            let cloudpickle = PyModule::import(py, "cloudpickle")?.getattr("dumps")?;
+            cloudpickle.call1((pyobj.clone_ref(py),))?
+        } else {
+            let pickle = PyModule::import(py, "pickle")?.getattr("dumps")?;
+            match pickle.call1((pyobj.clone_ref(py),)) {
+                Ok(dumped) => dumped,
+                Err(_) => {
+                    use_cloudpickle = true;
+                    let cloudpickle = PyModule::import(py, "cloudpickle")?.getattr("dumps")?;
+                    cloudpickle.call1((pyobj.clone_ref(py),))?
+                },
+            }
+        }
+        .extract::<PyBackedBytes>()
+    })?;
+
+    // Write pickle metadata
+    buf.push(use_cloudpickle as u8);
+    buf.extend_from_slice(&*PYTHON3_VERSION);
+
+    // Write UDF
+    buf.extend_from_slice(&dumped);
+    Ok(())
+}
+
+#[cfg(feature = "python")]
+pub fn python_object_deserialize(buf: &[u8]) -> PolarsResult<pyo3::Py<pyo3::PyAny>> {
+    use polars_error::polars_ensure;
+    use pyo3::Python;
+    use pyo3::types::{PyAnyMethods, PyBytes, PyModule};
+
+    use crate::python_function::PYTHON3_VERSION;
+
+    // Handle pickle metadata
+    let use_cloudpickle = buf[0] != 0;
+    if use_cloudpickle {
+        let ser_py_version = &buf[1..3];
+        let cur_py_version = *PYTHON3_VERSION;
+        polars_ensure!(
+            ser_py_version == cur_py_version,
+            InvalidOperation:
+            "current Python version {:?} does not match the Python version used to serialize the UDF {:?}",
+            (3, cur_py_version[0], cur_py_version[1]),
+            (3, ser_py_version[0], ser_py_version[1] )
+        );
+    }
+    let buf = &buf[3..];
+
+    Python::with_gil(|py| {
+        let loads = PyModule::import(py, "pickle")?.getattr("loads")?;
+        let arg = (PyBytes::new(py, buf),);
+        let python_function = loads.call1(arg)?;
+        Ok(python_function.into())
+    })
+}
+
 #[cfg(test)]
 mod tests {
     #[test]