polars_io/csv/write/
write_impl.rs

Help
1mod serializer;
2
3use std::io::Write;
4
5use arrow::array::NullArray;
6use arrow::legacy::time_zone::Tz;
7use polars_core::POOL;
8use polars_core::prelude::*;
9use polars_error::polars_ensure;
10use rayon::prelude::*;
11use serializer::{serializer_for, string_serializer};
12
13use crate::csv::write::SerializeOptions;
14
15pub(crate) fn write<W: Write>(
16    writer: &mut W,
17    df: &DataFrame,
18    chunk_size: usize,
19    options: &SerializeOptions,
20    n_threads: usize,
21) -> PolarsResult<()> {
22    for s in df.get_columns() {
23        let nested = match s.dtype() {
24            DataType::List(_) => true,
25            #[cfg(feature = "dtype-struct")]
26            DataType::Struct(_) => true,
27            #[cfg(feature = "object")]
28            DataType::Object(_) => {
29                return Err(PolarsError::ComputeError(
30                    "csv writer does not support object dtype".into(),
31                ));
32            },
33            _ => false,
34        };
35        polars_ensure!(
36            !nested,
37            ComputeError: "CSV format does not support nested data",
38        );
39    }
40
41    // Check that the double quote is valid UTF-8.
42    polars_ensure!(
43        std::str::from_utf8(&[options.quote_char, options.quote_char]).is_ok(),
44        ComputeError: "quote char results in invalid utf-8",
45    );
46
47    let (datetime_formats, time_zones): (Vec<&str>, Vec<Option<Tz>>) = df
48        .get_columns()
49        .iter()
50        .map(|column| match column.dtype() {
51            DataType::Datetime(TimeUnit::Milliseconds, tz) => {
52                let (format, tz_parsed) = match tz {
53                    #[cfg(feature = "timezones")]
54                    Some(tz) => (
55                        options
56                            .datetime_format
57                            .as_deref()
58                            .unwrap_or("%FT%H:%M:%S.%3f%z"),
59                        tz.parse::<Tz>().ok(),
60                    ),
61                    _ => (
62                        options
63                            .datetime_format
64                            .as_deref()
65                            .unwrap_or("%FT%H:%M:%S.%3f"),
66                        None,
67                    ),
68                };
69                (format, tz_parsed)
70            },
71            DataType::Datetime(TimeUnit::Microseconds, tz) => {
72                let (format, tz_parsed) = match tz {
73                    #[cfg(feature = "timezones")]
74                    Some(tz) => (
75                        options
76                            .datetime_format
77                            .as_deref()
78                            .unwrap_or("%FT%H:%M:%S.%6f%z"),
79                        tz.parse::<Tz>().ok(),
80                    ),
81                    _ => (
82                        options
83                            .datetime_format
84                            .as_deref()
85                            .unwrap_or("%FT%H:%M:%S.%6f"),
86                        None,
87                    ),
88                };
89                (format, tz_parsed)
90            },
91            DataType::Datetime(TimeUnit::Nanoseconds, tz) => {
92                let (format, tz_parsed) = match tz {
93                    #[cfg(feature = "timezones")]
94                    Some(tz) => (
95                        options
96                            .datetime_format
97                            .as_deref()
98                            .unwrap_or("%FT%H:%M:%S.%9f%z"),
99                        tz.parse::<Tz>().ok(),
100                    ),
101                    _ => (
102                        options
103                            .datetime_format
104                            .as_deref()
105                            .unwrap_or("%FT%H:%M:%S.%9f"),
106                        None,
107                    ),
108                };
109                (format, tz_parsed)
110            },
111            _ => ("", None),
112        })
113        .unzip();
114
115    let len = df.height();
116    let total_rows_per_pool_iter = n_threads * chunk_size;
117
118    let mut n_rows_finished = 0;
119
120    let mut buffers: Vec<_> = (0..n_threads).map(|_| (Vec::new(), Vec::new())).collect();
121    while n_rows_finished < len {
122        let buf_writer = |thread_no, write_buffer: &mut Vec<_>, serializers_vec: &mut Vec<_>| {
123            let thread_offset = thread_no * chunk_size;
124            let total_offset = n_rows_finished + thread_offset;
125            let mut df = df.slice(total_offset as i64, chunk_size);
126            // the `series.iter` needs rechunked series.
127            // we don't do this on the whole as this probably needs much less rechunking
128            // so will be faster.
129            // and allows writing `pl.concat([df] * 100, rechunk=False).write_csv()` as the rechunk
130            // would go OOM
131            df.as_single_chunk();
132            let cols = df.get_columns();
133
134            // SAFETY:
135            // the bck thinks the lifetime is bounded to write_buffer_pool, but at the time we return
136            // the vectors the buffer pool, the series have already been removed from the buffers
137            // in other words, the lifetime does not leave this scope
138            let cols = unsafe { std::mem::transmute::<&[Column], &[Column]>(cols) };
139
140            if df.is_empty() {
141                return Ok(());
142            }
143
144            if serializers_vec.is_empty() {
145                *serializers_vec = cols
146                    .iter()
147                    .enumerate()
148                    .map(|(i, col)| {
149                        serializer_for(
150                            &*col.as_materialized_series().chunks()[0],
151                            options,
152                            col.dtype(),
153                            datetime_formats[i],
154                            time_zones[i],
155                        )
156                    })
157                    .collect::<Result<_, _>>()?;
158            } else {
159                debug_assert_eq!(serializers_vec.len(), cols.len());
160                for (col_iter, col) in std::iter::zip(serializers_vec.iter_mut(), cols) {
161                    col_iter.update_array(&*col.as_materialized_series().chunks()[0]);
162                }
163            }
164
165            let serializers = serializers_vec.as_mut_slice();
166
167            let len = std::cmp::min(cols[0].len(), chunk_size);
168
169            for _ in 0..len {
170                serializers[0].serialize(write_buffer, options);
171                for serializer in &mut serializers[1..] {
172                    write_buffer.push(options.separator);
173                    serializer.serialize(write_buffer, options);
174                }
175
176                write_buffer.extend_from_slice(options.line_terminator.as_bytes());
177            }
178
179            Ok(())
180        };
181
182        if n_threads > 1 {
183            POOL.install(|| {
184                buffers
185                    .par_iter_mut()
186                    .enumerate()
187                    .map(|(i, (w, s))| buf_writer(i, w, s))
188                    .collect::<PolarsResult<()>>()
189            })?;
190        } else {
191            let (w, s) = &mut buffers[0];
192            buf_writer(0, w, s)?;
193        }
194
195        for (write_buffer, _) in &mut buffers {
196            writer.write_all(write_buffer)?;
197            write_buffer.clear();
198        }
199
200        n_rows_finished += total_rows_per_pool_iter;
201    }
202    Ok(())
203}
204
205/// Writes a CSV header to `writer`.
206pub(crate) fn write_header<W: Write>(
207    writer: &mut W,
208    names: &[&str],
209    options: &SerializeOptions,
210) -> PolarsResult<()> {
211    let mut header = Vec::new();
212
213    // A hack, but it works for this case.
214    let fake_arr = NullArray::new(ArrowDataType::Null, 0);
215    let mut names_serializer = string_serializer(
216        |iter: &mut std::slice::Iter<&str>| iter.next().copied(),
217        options,
218        |_| names.iter(),
219        &fake_arr,
220    );
221    for i in 0..names.len() {
222        names_serializer.serialize(&mut header, options);
223        if i != names.len() - 1 {
224            header.push(options.separator);
225        }
226    }
227    header.extend_from_slice(options.line_terminator.as_bytes());
228    writer.write_all(&header)?;
229    Ok(())
230}
231
232/// Writes a UTF-8 BOM to `writer`.
233pub(crate) fn write_bom<W: Write>(writer: &mut W) -> PolarsResult<()> {
234    const BOM: [u8; 3] = [0xEF, 0xBB, 0xBF];
235    writer.write_all(&BOM)?;
236    Ok(())
237}
polars_io/csv/write/write_impl.rs

polars_io/csv/write/
write_impl.rs
