polars_python/dataframe/
serde.rs

1use std::io::{BufReader, BufWriter};
2use std::ops::Deref;
3
4use polars::prelude::*;
5use polars_io::mmap::ReaderBytes;
6use pyo3::prelude::*;
7
8use super::PyDataFrame;
9use crate::exceptions::ComputeError;
10use crate::file::{get_file_like, get_mmap_bytes_reader};
11use crate::utils::EnterPolarsExt;
12
13#[pymethods]
14impl PyDataFrame {
15    /// Serialize into binary data.
16    fn serialize_binary(&mut self, py: Python, py_f: PyObject) -> PyResult<()> {
17        let file = get_file_like(py_f, true)?;
18        let mut writer = BufWriter::new(file);
19
20        py.enter_polars(|| self.df.serialize_into_writer(&mut writer))
21    }
22
23    /// Deserialize a file-like object containing binary data into a DataFrame.
24    #[staticmethod]
25    fn deserialize_binary(py: Python, py_f: PyObject) -> PyResult<Self> {
26        let file = get_file_like(py_f, false)?;
27        let mut file = BufReader::new(file);
28
29        py.enter_polars_df(|| DataFrame::deserialize_from_reader(&mut file))
30    }
31
32    /// Serialize into a JSON string.
33    #[cfg(feature = "json")]
34    pub fn serialize_json(&mut self, py: Python, py_f: PyObject) -> PyResult<()> {
35        let file = get_file_like(py_f, true)?;
36        let writer = BufWriter::new(file);
37        py.enter_polars(|| {
38            serde_json::to_writer(writer, &self.df)
39                .map_err(|err| ComputeError::new_err(err.to_string()))
40        })
41    }
42
43    /// Deserialize a file-like object containing JSON string data into a DataFrame.
44    #[staticmethod]
45    #[cfg(feature = "json")]
46    pub fn deserialize_json(py: Python, py_f: Bound<PyAny>) -> PyResult<Self> {
47        let mut mmap_bytes_r = get_mmap_bytes_reader(&py_f)?;
48
49        py.enter_polars(move || {
50            let mmap_read: ReaderBytes = (&mut mmap_bytes_r).into();
51            let bytes = mmap_read.deref();
52            let df = serde_json::from_slice::<DataFrame>(bytes)
53                .map_err(|err| ComputeError::new_err(err.to_string()))?;
54            PyResult::Ok(df.into())
55        })
56    }
57}