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::error::PyPolarsErr;
10use crate::exceptions::ComputeError;
11use crate::file::{get_file_like, get_mmap_bytes_reader};
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.allow_threads(|| {
21            self.df
22                .serialize_into_writer(&mut writer)
23                .map_err(|e| PyPolarsErr::from(e).into())
24        })
25    }
26
27    /// Deserialize a file-like object containing binary data into a DataFrame.
28    #[staticmethod]
29    fn deserialize_binary(py: Python, py_f: PyObject) -> PyResult<Self> {
30        let file = get_file_like(py_f, false)?;
31        let mut file = BufReader::new(file);
32
33        py.allow_threads(|| {
34            DataFrame::deserialize_from_reader(&mut file)
35                .map_err(|e| PyPolarsErr::from(e).into())
36                .map(|x| x.into())
37        })
38    }
39
40    /// Serialize into a JSON string.
41    #[cfg(feature = "json")]
42    pub fn serialize_json(&mut self, py: Python, py_f: PyObject) -> PyResult<()> {
43        let file = get_file_like(py_f, true)?;
44        let writer = BufWriter::new(file);
45        py.allow_threads(|| {
46            serde_json::to_writer(writer, &self.df)
47                .map_err(|err| ComputeError::new_err(err.to_string()))
48        })
49    }
50
51    /// Deserialize a file-like object containing JSON string data into a DataFrame.
52    #[staticmethod]
53    #[cfg(feature = "json")]
54    pub fn deserialize_json(py: Python, py_f: Bound<PyAny>) -> PyResult<Self> {
55        let mut mmap_bytes_r = get_mmap_bytes_reader(&py_f)?;
56
57        py.allow_threads(move || {
58            let mmap_read: ReaderBytes = (&mut mmap_bytes_r).into();
59            let bytes = mmap_read.deref();
60            let df = serde_json::from_slice::<DataFrame>(bytes)
61                .map_err(|err| ComputeError::new_err(err.to_string()))?;
62            Ok(df.into())
63        })
64    }
65}