polars_python/expr/
serde.rs

1use std::io::{BufReader, BufWriter};
2
3use polars::lazy::prelude::Expr;
4use polars_utils::pl_serialize;
5use pyo3::prelude::*;
6use pyo3::pybacked::PyBackedBytes;
7use pyo3::types::PyBytes;
8
9use crate::PyExpr;
10use crate::error::PyPolarsErr;
11use crate::exceptions::ComputeError;
12use crate::file::get_file_like;
13
14#[pymethods]
15impl PyExpr {
16    // Pickle we set FC is false, as that is used for caching (compact is faster) and is not intended to be used
17    // across different versions.
18    fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
19        // Used in pickle/pickling
20        let mut writer: Vec<u8> = vec![];
21        pl_serialize::SerializeOptions::default()
22            .serialize_into_writer::<_, _, false>(&mut writer, &self.inner)
23            .map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
24
25        Ok(PyBytes::new(py, &writer))
26    }
27
28    fn __setstate__(&mut self, state: &Bound<PyAny>) -> PyResult<()> {
29        // Used in pickle/pickling
30        let bytes = state.extract::<PyBackedBytes>()?;
31        self.inner = pl_serialize::SerializeOptions::default()
32            .deserialize_from_reader::<_, _, false>(&*bytes)
33            .map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
34        Ok(())
35    }
36
37    /// Serialize into binary data.
38    fn serialize_binary(&self, py_f: PyObject) -> PyResult<()> {
39        let file = get_file_like(py_f, true)?;
40        let writer = BufWriter::new(file);
41        pl_serialize::SerializeOptions::default()
42            .serialize_into_writer::<_, _, true>(writer, &self.inner)
43            .map_err(|err| ComputeError::new_err(err.to_string()))
44    }
45
46    /// Serialize into a JSON string.
47    #[cfg(feature = "json")]
48    fn serialize_json(&self, py_f: PyObject) -> PyResult<()> {
49        let file = get_file_like(py_f, true)?;
50        let writer = BufWriter::new(file);
51        serde_json::to_writer(writer, &self.inner)
52            .map_err(|err| ComputeError::new_err(err.to_string()))
53    }
54
55    /// Deserialize a file-like object containing binary data into an Expr.
56    #[staticmethod]
57    fn deserialize_binary(py_f: PyObject) -> PyResult<PyExpr> {
58        let file = get_file_like(py_f, false)?;
59        let reader = BufReader::new(file);
60        let expr: Expr = pl_serialize::SerializeOptions::default()
61            .deserialize_from_reader::<_, _, true>(reader)
62            .map_err(|err| ComputeError::new_err(err.to_string()))?;
63        Ok(expr.into())
64    }
65
66    /// Deserialize a file-like object containing JSON string data into an Expr.
67    #[staticmethod]
68    #[cfg(feature = "json")]
69    fn deserialize_json(py_f: PyObject) -> PyResult<PyExpr> {
70        // it is faster to first read to memory and then parse: https://github.com/serde-rs/json/issues/160
71        // so don't bother with files.
72        let mut json = String::new();
73        let _ = get_file_like(py_f, false)?
74            .read_to_string(&mut json)
75            .unwrap();
76
77        // SAFETY:
78        // We skipped the serializing/deserializing of the static in lifetime in `DataType`
79        // so we actually don't have a lifetime at all when serializing.
80
81        // &str still has a lifetime. But it's ok, because we drop it immediately
82        // in this scope.
83        let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
84
85        let inner: Expr = serde_json::from_str(json).map_err(|_| {
86            let msg = "could not deserialize input into an expression";
87            ComputeError::new_err(msg)
88        })?;
89        Ok(inner.into())
90    }
91}
polars_python/expr/serde.rs

polars_python/expr/
serde.rs
