polars_core/serde/
df.rs

1use std::io::{Read, Seek};
2use std::sync::Arc;
3
4use polars_arrow::datatypes::Metadata;
5use polars_arrow::io::ipc::read::{StreamReader, StreamState, read_stream_metadata};
6use polars_arrow::io::ipc::write::WriteOptions;
7use polars_error::{PolarsResult, polars_err, to_compute_err};
8use polars_utils::format_pl_smallstr;
9use polars_utils::pl_serialize::deserialize_map_bytes;
10use polars_utils::pl_str::PlSmallStr;
11use serde::de::Error;
12use serde::*;
13
14use crate::chunked_array::flags::StatisticsFlags;
15use crate::config;
16use crate::frame::chunk_df_for_writing;
17use crate::prelude::{CompatLevel, DataFrame, SchemaExt};
18use crate::schema::Schema;
19use crate::utils::accumulate_dataframes_vertical_unchecked;
20
21const FLAGS_KEY: PlSmallStr = PlSmallStr::from_static("_PL_FLAGS");
22
23impl DataFrame {
24    pub fn serialize_into_json(&self, writer: &mut dyn std::io::Write) -> PolarsResult<()> {
25        serde_json::to_writer(writer, self).map_err(to_compute_err)
26    }
27
28    pub fn deserialize_from_json(json: &[u8]) -> PolarsResult<Self> {
29        serde_json::from_slice(json).map_err(to_compute_err)
30    }
31
32    pub fn serialize_into_writer(&mut self, writer: &mut dyn std::io::Write) -> PolarsResult<()> {
33        let schema = self.schema();
34
35        if schema.iter_values().any(|x| x.is_object()) {
36            return Err(polars_err!(
37                ComputeError:
38                "serializing data of type Object is not supported",
39            ));
40        }
41
42        let mut ipc_writer = polars_arrow::io::ipc::write::StreamWriter::new(
43            writer,
44            WriteOptions { compression: None },
45        );
46
47        ipc_writer.set_custom_schema_metadata(Arc::new(Metadata::from_iter(
48            self.columns().iter().map(|c| {
49                (
50                    format_pl_smallstr!("{}{}", FLAGS_KEY, c.name()),
51                    PlSmallStr::from(c.get_flags().bits().to_string()),
52                )
53            }),
54        )));
55
56        ipc_writer.set_custom_schema_metadata(Arc::new(Metadata::from([(
57            FLAGS_KEY,
58            serde_json::to_string(
59                &self
60                    .columns()
61                    .iter()
62                    .map(|s| s.get_flags().bits())
63                    .collect::<Vec<u32>>(),
64            )
65            .map_err(to_compute_err)?
66            .into(),
67        )])));
68
69        ipc_writer.start(&schema.to_arrow(CompatLevel::newest()), None)?;
70
71        for batch in chunk_df_for_writing(self, 512 * 512)?.iter_chunks(CompatLevel::newest(), true)
72        {
73            ipc_writer.write(&batch, None)?;
74        }
75
76        ipc_writer.finish()?;
77
78        Ok(())
79    }
80
81    pub fn deserialize_from_reader<T: Read + Seek>(reader: &mut T) -> PolarsResult<Self> {
82        let mut md = read_stream_metadata(reader)?;
83        let pl_schema = Schema::from_arrow_schema(&md.schema);
84
85        let custom_metadata = md.custom_schema_metadata.take();
86
87        let reader = StreamReader::new(reader, md, None);
88        let dfs = reader
89            .into_iter()
90            .map_while(|batch| match batch {
91                Ok(StreamState::Some(batch)) => Some(Ok(DataFrame::from(batch))),
92                Ok(StreamState::Waiting) => None,
93                Err(e) => Some(Err(e)),
94            })
95            .collect::<PolarsResult<Vec<DataFrame>>>()?;
96
97        if dfs.is_empty() {
98            return Ok(DataFrame::empty_with_schema(&pl_schema));
99        }
100
101        let mut df = accumulate_dataframes_vertical_unchecked(dfs);
102
103        // Set custom metadata (fallible)
104        (|| {
105            let custom_metadata = custom_metadata?;
106            let flags = custom_metadata.get(&FLAGS_KEY)?;
107
108            let flags: PolarsResult<Vec<u32>> = serde_json::from_str(flags).map_err(to_compute_err);
109
110            let verbose = config::verbose();
111
112            if let Err(e) = &flags {
113                if verbose {
114                    eprintln!("DataFrame::read_ipc: Error parsing metadata flags: {e}");
115                }
116            }
117
118            let flags = flags.ok()?;
119
120            if flags.len() != df.width() {
121                if verbose {
122                    eprintln!(
123                        "DataFrame::read_ipc: Metadata flags width mismatch: {} != {}",
124                        flags.len(),
125                        df.width()
126                    );
127                }
128
129                return None;
130            }
131
132            let mut n_set = 0;
133
134            for (c, v) in unsafe { df.columns_mut_retain_schema() }
135                .iter_mut()
136                .zip(flags)
137            {
138                if let Some(flags) = StatisticsFlags::from_bits(v) {
139                    n_set += c.set_flags(flags) as usize;
140                }
141            }
142
143            if verbose {
144                eprintln!(
145                    "DataFrame::read_ipc: Loaded metadata for {} / {} columns",
146                    n_set,
147                    df.width()
148                );
149            }
150
151            Some(())
152        })();
153
154        Ok(df)
155    }
156}
157
158impl Serialize for DataFrame {
159    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
160    where
161        S: Serializer,
162    {
163        use serde::ser::Error;
164
165        let mut bytes = vec![];
166        self.clone()
167            .serialize_into_writer(&mut bytes)
168            .map_err(S::Error::custom)?;
169
170        serializer.serialize_bytes(bytes.as_slice())
171    }
172}
173
174impl<'de> Deserialize<'de> for DataFrame {
175    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
176    where
177        D: Deserializer<'de>,
178    {
179        deserialize_map_bytes(deserializer, |b| {
180            let v = &mut b.as_ref();
181            let mut reader = std::io::Cursor::new(v);
182            Self::deserialize_from_reader(&mut reader)
183        })?
184        .map_err(D::Error::custom)
185    }
186}
187
188#[cfg(feature = "dsl-schema")]
189impl schemars::JsonSchema for DataFrame {
190    fn schema_name() -> std::borrow::Cow<'static, str> {
191        "DataFrame".into()
192    }
193
194    fn schema_id() -> std::borrow::Cow<'static, str> {
195        std::borrow::Cow::Borrowed(concat!(module_path!(), "::", "DataFrame"))
196    }
197
198    fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
199        Vec::<u8>::json_schema(generator)
200    }
201}
polars_core/serde/df.rs

polars_core/serde/
df.rs
