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