polars_python/lazyframe/
serde.rs

1use std::io::{BufReader, BufWriter, Read};
2
3use pyo3::prelude::*;
4
5use super::PyLazyFrame;
6use crate::exceptions::ComputeError;
7use crate::file::get_file_like;
8use crate::prelude::*;
9use crate::utils::EnterPolarsExt;
10
11#[pymethods]
12#[allow(clippy::should_implement_trait)]
13impl PyLazyFrame {
14    /// Serialize into binary data.
15    fn serialize_binary(&self, py: Python<'_>, py_f: PyObject) -> PyResult<()> {
16        let file = get_file_like(py_f, true)?;
17        let writer = BufWriter::new(file);
18        py.enter_polars(|| self.ldf.logical_plan.serialize_versioned(writer))
19    }
20
21    /// Serialize into a JSON string.
22    #[cfg(feature = "json")]
23    fn serialize_json(&self, py: Python<'_>, py_f: PyObject) -> PyResult<()> {
24        let file = get_file_like(py_f, true)?;
25        let writer = BufWriter::new(file);
26        py.enter_polars(|| {
27            serde_json::to_writer(writer, &self.ldf.logical_plan)
28                .map_err(|err| ComputeError::new_err(err.to_string()))
29        })
30    }
31
32    /// Deserialize a file-like object containing binary data into a LazyFrame.
33    #[staticmethod]
34    fn deserialize_binary(py: Python<'_>, py_f: PyObject) -> PyResult<Self> {
35        let file = get_file_like(py_f, false)?;
36        let reader = BufReader::new(file);
37
38        let lp: DslPlan = py.enter_polars(|| DslPlan::deserialize_versioned(reader))?;
39        Ok(LazyFrame::from(lp).into())
40    }
41
42    /// Deserialize a file-like object containing JSON string data into a LazyFrame.
43    #[staticmethod]
44    #[cfg(feature = "json")]
45    fn deserialize_json(py: Python<'_>, py_f: PyObject) -> PyResult<Self> {
46        // it is faster to first read to memory and then parse: https://github.com/serde-rs/json/issues/160
47        // so don't bother with files.
48        let mut json = String::new();
49        get_file_like(py_f, false)?
50            .read_to_string(&mut json)
51            .unwrap();
52
53        // SAFETY:
54        // We skipped the serializing/deserializing of the static in lifetime in `DataType`
55        // so we actually don't have a lifetime at all when serializing.
56
57        // &str still has a lifetime. But it's ok, because we drop it immediately
58        // in this scope.
59        let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
60
61        let lp = py.enter_polars(|| {
62            serde_json::from_str::<DslPlan>(json)
63                .map_err(|err| ComputeError::new_err(err.to_string()))
64        })?;
65        Ok(LazyFrame::from(lp).into())
66    }
67}