polars_utils/
pl_serialize.rs

1//! Centralized Polars serialization entry.
2//!
3//! Currently provides two serialization scheme's.
4//! - Self-describing (and thus more forward compatible) activated with `FC: true`
5//! - Compact activated with `FC: false`
6use std::io::{Read, Write};
7
8use polars_error::{PolarsResult, to_compute_err};
9
10fn config() -> bincode::config::Configuration {
11    bincode::config::standard()
12        .with_no_limit()
13        .with_variable_int_encoding()
14}
15
16fn serialize_impl<T, const FC: bool>(mut writer: &mut dyn Write, value: &T) -> PolarsResult<()>
17where
18    T: serde::ser::Serialize,
19{
20    if FC {
21        serialize_dsl(writer, value)
22    } else {
23        bincode::serde::encode_into_std_write(value, &mut writer, config())
24            .map_err(to_compute_err)
25            .map(|_| ())
26    }
27}
28
29fn deserialize_impl<T, const FC: bool>(mut reader: &mut dyn Read) -> PolarsResult<T>
30where
31    T: serde::de::DeserializeOwned,
32{
33    if FC {
34        deserialize_dsl(reader)
35    } else {
36        bincode::serde::decode_from_std_read(&mut reader, config()).map_err(to_compute_err)
37    }
38}
39
40/// Mainly used to enable compression when serializing the final outer value.
41/// For intermediate serialization steps, the function in the module should
42/// be used instead.
43pub struct SerializeOptions {
44    compression: bool,
45}
46
47impl SerializeOptions {
48    pub fn with_compression(mut self, compression: bool) -> Self {
49        self.compression = compression;
50        self
51    }
52
53    pub fn serialize_into_writer<W, T, const FC: bool>(
54        &self,
55        mut writer: W,
56        value: &T,
57    ) -> PolarsResult<()>
58    where
59        W: std::io::Write,
60        T: serde::ser::Serialize,
61    {
62        if self.compression {
63            let mut compr_writer =
64                flate2::write::ZlibEncoder::new(writer, flate2::Compression::fast());
65            serialize_impl::<_, FC>(&mut compr_writer, value)
66        } else {
67            serialize_impl::<_, FC>(&mut writer, value)
68        }
69    }
70
71    pub fn deserialize_from_reader<T, R, const FC: bool>(&self, mut reader: R) -> PolarsResult<T>
72    where
73        T: serde::de::DeserializeOwned,
74        R: std::io::Read,
75    {
76        if self.compression {
77            let mut compr_reader = flate2::read::ZlibDecoder::new(reader);
78            deserialize_impl::<_, FC>(&mut compr_reader)
79        } else {
80            deserialize_impl::<_, FC>(&mut reader)
81        }
82    }
83
84    pub fn serialize_to_bytes<T, const FC: bool>(&self, value: &T) -> PolarsResult<Vec<u8>>
85    where
86        T: serde::ser::Serialize,
87    {
88        let mut v = vec![];
89
90        self.serialize_into_writer::<_, _, FC>(&mut v, value)?;
91
92        Ok(v)
93    }
94}
95
96#[allow(clippy::derivable_impls)]
97impl Default for SerializeOptions {
98    fn default() -> Self {
99        Self { compression: false }
100    }
101}
102
103pub fn serialize_into_writer<W, T, const FC: bool>(mut writer: W, value: &T) -> PolarsResult<()>
104where
105    W: std::io::Write,
106    T: serde::ser::Serialize,
107{
108    serialize_impl::<_, FC>(&mut writer, value)
109}
110
111pub fn deserialize_from_reader<T, R, const FC: bool>(mut reader: R) -> PolarsResult<T>
112where
113    T: serde::de::DeserializeOwned,
114    R: std::io::Read,
115{
116    deserialize_impl::<_, FC>(&mut reader)
117}
118
119pub fn serialize_to_bytes<T, const FC: bool>(value: &T) -> PolarsResult<Vec<u8>>
120where
121    T: serde::ser::Serialize,
122{
123    let mut v = vec![];
124    serialize_into_writer::<_, _, FC>(&mut v, value)?;
125    Ok(v)
126}
127
128/// Serialize function customized for `DslPlan`, with stack overflow protection.
129///
130/// This is also the `FC = true` format, so that every type is only instantiated for a single
131/// serializer.
132pub fn serialize_dsl<T>(writer: &mut dyn Write, value: &T) -> PolarsResult<()>
133where
134    T: serde::ser::Serialize,
135{
136    let mut s = rmp_serde::Serializer::new(writer).with_struct_map();
137    let s = serde_stacker::Serializer::new(&mut s);
138    value.serialize(s).map_err(to_compute_err)
139}
140
141/// Deserialize function customized for `DslPlan`, with stack overflow protection.
142pub fn deserialize_dsl<T>(reader: &mut dyn Read) -> PolarsResult<T>
143where
144    T: serde::de::DeserializeOwned,
145{
146    let mut de = rmp_serde::Deserializer::new(reader);
147    de.set_max_depth(usize::MAX);
148    let de = serde_stacker::Deserializer::new(&mut de);
149    T::deserialize(de).map_err(to_compute_err)
150}
151
152/// Potentially avoids copying memory compared to a naive `Vec::<u8>::deserialize`.
153///
154/// This is essentially boilerplate for visiting bytes without copying where possible.
155pub fn deserialize_map_bytes<'de, D, O>(
156    deserializer: D,
157    mut func: impl for<'b> FnMut(std::borrow::Cow<'b, [u8]>) -> O,
158) -> Result<O, D::Error>
159where
160    D: serde::de::Deserializer<'de>,
161{
162    // Lets us avoid monomorphizing the visitor
163    let mut out: Option<O> = None;
164    struct V<'f>(&'f mut dyn for<'b> FnMut(std::borrow::Cow<'b, [u8]>));
165
166    deserializer.deserialize_bytes(V(&mut |v| drop(out.replace(func(v)))))?;
167
168    return Ok(out.unwrap());
169
170    impl<'de> serde::de::Visitor<'de> for V<'_> {
171        type Value = ();
172
173        fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
174            formatter.write_str("deserialize_map_bytes")
175        }
176
177        fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
178        where
179            E: serde::de::Error,
180        {
181            self.0(std::borrow::Cow::Borrowed(v));
182            Ok(())
183        }
184
185        fn visit_byte_buf<E>(self, v: Vec<u8>) -> Result<Self::Value, E>
186        where
187            E: serde::de::Error,
188        {
189            self.0(std::borrow::Cow::Owned(v));
190            Ok(())
191        }
192
193        fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
194        where
195            A: serde::de::SeqAccess<'de>,
196        {
197            // This is not ideal, but we hit here if the serialization format is JSON.
198            let bytes = std::iter::from_fn(|| seq.next_element::<u8>().transpose())
199                .collect::<Result<Vec<_>, A::Error>>()?;
200
201            self.0(std::borrow::Cow::Owned(bytes));
202            Ok(())
203        }
204    }
205}
206
207thread_local! {
208    pub static USE_CLOUDPICKLE: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
209}
210
211#[cfg(feature = "python")]
212pub fn python_object_serialize(
213    pyobj: &pyo3::Py<pyo3::PyAny>,
214    buf: &mut Vec<u8>,
215) -> PolarsResult<()> {
216    use pyo3::Python;
217    use pyo3::pybacked::PyBackedBytes;
218    use pyo3::types::{PyAnyMethods, PyModule};
219
220    use crate::python_function::PYTHON3_VERSION;
221
222    let mut use_cloudpickle = USE_CLOUDPICKLE.get();
223    let dumped = Python::attach(|py| {
224        // Pickle with whatever pickling method was selected.
225        if use_cloudpickle {
226            let cloudpickle = PyModule::import(py, "cloudpickle")?.getattr("dumps")?;
227            cloudpickle.call1((pyobj.clone_ref(py),))?
228        } else {
229            let pickle = PyModule::import(py, "pickle")?.getattr("dumps")?;
230            match pickle.call1((pyobj.clone_ref(py),)) {
231                Ok(dumped) => dumped,
232                Err(_) => {
233                    use_cloudpickle = true;
234                    let cloudpickle = PyModule::import(py, "cloudpickle")?.getattr("dumps")?;
235                    cloudpickle.call1((pyobj.clone_ref(py),))?
236                },
237            }
238        }
239        .extract::<PyBackedBytes>()
240        .map_err(pyo3::PyErr::from)
241    })?;
242
243    // Write pickle metadata
244    buf.push(use_cloudpickle as u8);
245    buf.extend_from_slice(&*PYTHON3_VERSION);
246
247    // Write UDF
248    buf.extend_from_slice(&dumped);
249    Ok(())
250}
251
252#[cfg(feature = "python")]
253pub fn python_object_deserialize(buf: &[u8]) -> PolarsResult<pyo3::Py<pyo3::PyAny>> {
254    use polars_error::polars_ensure;
255    use pyo3::Python;
256    use pyo3::types::{PyAnyMethods, PyBytes, PyModule};
257
258    use crate::python_function::PYTHON3_VERSION;
259
260    // Handle pickle metadata
261    let use_cloudpickle = buf[0] != 0;
262    if use_cloudpickle {
263        let ser_py_version = &buf[1..3];
264        let cur_py_version = *PYTHON3_VERSION;
265        polars_ensure!(
266            ser_py_version == cur_py_version,
267            InvalidOperation:
268            "current Python version {:?} does not match the Python version used to serialize the UDF {:?}",
269            (3, cur_py_version[0], cur_py_version[1]),
270            (3, ser_py_version[0], ser_py_version[1] )
271        );
272    }
273    let buf = &buf[3..];
274
275    Python::attach(|py| {
276        let loads = PyModule::import(py, "pickle")?.getattr("loads")?;
277        let arg = (PyBytes::new(py, buf),);
278        let python_function = loads.call1(arg)?;
279        Ok(python_function.into())
280    })
281}
282
283#[cfg(test)]
284mod tests {
285    #[test]
286    fn test_serde_skip_enum() {
287        #[derive(Default, Debug, PartialEq)]
288        struct MyType(Option<usize>);
289
290        // Note: serde(skip) must be at the end of enums
291        #[derive(Debug, PartialEq, serde::Serialize, serde::Deserialize)]
292        enum Enum {
293            A,
294            #[serde(skip)]
295            B(MyType),
296        }
297
298        impl Default for Enum {
299            fn default() -> Self {
300                Self::B(MyType(None))
301            }
302        }
303
304        let v = Enum::A;
305        let b = super::serialize_to_bytes::<_, false>(&v).unwrap();
306        let r: Enum = super::deserialize_from_reader::<_, _, false>(b.as_slice()).unwrap();
307
308        assert_eq!(r, v);
309
310        let v = Enum::A;
311        let b = super::SerializeOptions::default()
312            .serialize_to_bytes::<_, false>(&v)
313            .unwrap();
314        let r: Enum = super::SerializeOptions::default()
315            .deserialize_from_reader::<_, _, false>(b.as_slice())
316            .unwrap();
317
318        assert_eq!(r, v);
319    }
320}
polars_utils/pl_serialize.rs

polars_utils/
pl_serialize.rs
