feat: Check Python version when deserializing UDFs

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

@@ -25,6 +25,7 @@ pub static mut CALL_COLUMNS_UDF_PYTHON: Option<
 pub static mut CALL_DF_UDF_PYTHON: Option<
     fn(s: DataFrame, lambda: &PyObject) -> PolarsResult<DataFrame>,
 > = None;
+#[cfg(feature = "serde")]
 pub(super) const MAGIC_BYTE_MARK: &[u8] = "PLPYUDF".as_bytes();
 
 #[derive(Clone, Debug)]
@@ -60,7 +61,7 @@ impl Serialize for PythonFunction {
         Python::with_gil(|py| {
             let pickle = PyModule::import_bound(py, "cloudpickle")
                 .or_else(|_| PyModule::import_bound(py, "pickle"))
-                .expect("Unable to import 'cloudpickle' or 'pickle'")
+                .expect("unable to import 'cloudpickle' or 'pickle'")
                 .getattr("dumps")
                 .unwrap();
 
@@ -86,9 +87,8 @@ impl<'a> Deserialize<'a> for PythonFunction {
         let bytes = Vec::<u8>::deserialize(deserializer)?;
 
         Python::with_gil(|py| {
-            let pickle = PyModule::import_bound(py, "cloudpickle")
-                .or_else(|_| PyModule::import_bound(py, "pickle"))
-                .expect("Unable to import 'pickle'")
+            let pickle = PyModule::import_bound(py, "pickle")
+                .expect("unable to import 'pickle'")
                 .getattr("loads")
                 .unwrap();
             let arg = (PyBytes::new_bound(py, &bytes),);
@@ -125,19 +125,36 @@ impl PythonUdfExpression {
 
     #[cfg(feature = "serde")]
     pub(crate) fn try_deserialize(buf: &[u8]) -> PolarsResult<Arc<dyn ColumnsUdf>> {
+        // Handle byte mark
         debug_assert!(buf.starts_with(MAGIC_BYTE_MARK));
-        // skip header
         let buf = &buf[MAGIC_BYTE_MARK.len()..];
+
+        // Handle pickle metadata
+        let use_cloudpickle = buf[0];
+        if use_cloudpickle != 0 {
+            let ser_py_version = buf[1];
+            let cur_py_version = get_python_minor_version();
+            polars_ensure!(
+                ser_py_version == cur_py_version,
+                InvalidOperation:
+                "current Python version (3.{}) does not match the Python version used to serialize the UDF (3.{})",
+                cur_py_version,
+                ser_py_version
+            );
+        }
+        let buf = &buf[2..];
+
+        // Load UDF metadata
         let mut reader = Cursor::new(buf);
         let (output_type, is_elementwise, returns_scalar): (Option<DataType>, bool, bool) =
             ciborium::de::from_reader(&mut reader).map_err(map_err)?;
 
         let remainder = &buf[reader.position() as usize..];
 
+        // Load UDF
         Python::with_gil(|py| {
-            let pickle = PyModule::import_bound(py, "cloudpickle")
-                .or_else(|_| PyModule::import_bound(py, "pickle"))
-                .expect("Unable to import 'pickle'")
+            let pickle = PyModule::import_bound(py, "pickle")
+                .expect("unable to import 'pickle'")
                 .getattr("loads")
                 .unwrap();
             let arg = (PyBytes::new_bound(py, remainder),);
@@ -189,26 +206,45 @@ impl ColumnsUdf for PythonUdfExpression {
 
     #[cfg(feature = "serde")]
     fn try_serialize(&self, buf: &mut Vec<u8>) -> PolarsResult<()> {
+        // Write byte marks
         buf.extend_from_slice(MAGIC_BYTE_MARK);
-        ciborium::ser::into_writer(
-            &(
-                self.output_type.clone(),
-                self.is_elementwise,
-                self.returns_scalar,
-            ),
-            &mut *buf,
-        )
-        .unwrap();
 
         Python::with_gil(|py| {
-            let pickle = PyModule::import_bound(py, "cloudpickle")
-                .or_else(|_| PyModule::import_bound(py, "pickle"))
-                .expect("Unable to import 'pickle'")
+            // Try pickle to serialize the UDF, otherwise fall back to cloudpickle.
+            let pickle = PyModule::import_bound(py, "pickle")
+                .expect("unable to import 'pickle'")
                 .getattr("dumps")
                 .unwrap();
-            let dumped = pickle
-                .call1((self.python_function.clone(),))
-                .map_err(from_pyerr)?;
+            let pickle_result = pickle.call1((self.python_function.clone(),));
+            let (dumped, use_cloudpickle, py_version) = match pickle_result {
+                Ok(dumped) => (dumped, false, 0),
+                Err(_) => {
+                    let cloudpickle = PyModule::import_bound(py, "cloudpickle")
+                        .map_err(from_pyerr)?
+                        .getattr("dumps")
+                        .unwrap();
+                    let dumped = cloudpickle
+                        .call1((self.python_function.clone(),))
+                        .map_err(from_pyerr)?;
+                    (dumped, true, get_python_minor_version())
+                },
+            };
+
+            // Write pickle metadata
+            buf.extend_from_slice(&[use_cloudpickle as u8, py_version]);
+
+            // Write UDF metadata
+            ciborium::ser::into_writer(
+                &(
+                    self.output_type.clone(),
+                    self.is_elementwise,
+                    self.returns_scalar,
+                ),
+                &mut *buf,
+            )
+            .unwrap();
+
+            // Write UDF
             let dumped = dumped.extract::<PyBackedBytes>().unwrap();
             buf.extend_from_slice(&dumped);
             Ok(())
@@ -298,3 +334,17 @@ impl Expr {
         }
     }
 }
+
+/// Get the minor Python version from the `sys` module.
+fn get_python_minor_version() -> u8 {
+    Python::with_gil(|py| {
+        PyModule::import_bound(py, "sys")
+            .unwrap()
+            .getattr("version_info")
+            .unwrap()
+            .getattr("minor")
+            .unwrap()
+            .extract()
+            .unwrap()
+    })
+}

py-polars/tests/unit/lazyframe/test_serde.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import io
-from typing import TYPE_CHECKING
+import sys
+from typing import TYPE_CHECKING, NamedTuple
 
 import pytest
 from hypothesis import example, given
 
 import polars as pl
-from polars.exceptions import ComputeError
+from polars.exceptions import ComputeError, InvalidOperationError
 from polars.testing import assert_frame_equal
 from polars.testing.parametric import dataframes
 
@@ -116,3 +117,60 @@ def test_lf_serde_scan(tmp_path: Path) -> None:
     result = pl.LazyFrame.deserialize(io.BytesIO(ser))
     assert_frame_equal(result, lf)
     assert_frame_equal(result.collect(), df)
+
+
+class MockVersionInfo(NamedTuple):
+    """Version info with different minor version."""
+
+    major: int = sys.version_info.major
+    minor: int = sys.version_info.minor + 1
+    micro: int = sys.version_info.micro
+
+
+@pytest.mark.filterwarnings("ignore::polars.exceptions.PolarsInefficientMapWarning")
+def test_lf_serde_lambda_version_specific(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    lf = pl.LazyFrame({"a": [1, 2, 3]}).select(
+        pl.col("a").map_elements(lambda x: x + 1, return_dtype=pl.Int64)
+    )
+    ser = lf.serialize()
+
+    # Same version
+    result = pl.LazyFrame.deserialize(io.BytesIO(ser))
+    expected = pl.LazyFrame({"a": [2, 3, 4]})
+    assert_frame_equal(result, expected)
+
+    # Different version
+    monkeypatch.setattr(sys, "version_info", MockVersionInfo())
+
+    with pytest.raises(
+        InvalidOperationError,
+        match="does not match the Python version used to serialize the UDF",
+    ):
+        pl.LazyFrame.deserialize(io.BytesIO(ser)).collect()
+
+
+def custom_function(x: pl.Series) -> pl.Series:
+    return x + 1
+
+
+@pytest.mark.filterwarnings("ignore::polars.exceptions.PolarsInefficientMapWarning")
+def test_lf_serde_function_version_specific(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    lf = pl.LazyFrame({"a": [1, 2, 3]}).select(
+        pl.col("a").map_batches(custom_function, return_dtype=pl.Int64)
+    )
+    ser = lf.serialize()
+
+    # Same version
+    result = pl.LazyFrame.deserialize(io.BytesIO(ser))
+    expected = pl.LazyFrame({"a": [2, 3, 4]})
+    assert_frame_equal(result, expected)
+
+    # Different version
+    monkeypatch.setattr(sys, "version_info", MockVersionInfo())
+
+    result = pl.LazyFrame.deserialize(io.BytesIO(ser))
+    assert_frame_equal(result, expected)