polars_io/avro/
write.rs

1use std::io::Write;
2
3pub use Compression as AvroCompression;
4pub use polars_arrow::io::avro::avro_schema::file::Compression;
5use polars_arrow::io::avro::avro_schema::{self};
6use polars_arrow::io::avro::write;
7use polars_core::error::to_compute_err;
8use polars_core::prelude::*;
9
10use crate::shared::{SerWriter, schema_to_arrow_checked};
11
12/// Write a [`DataFrame`] to [Apache Avro] format
13///
14/// [Apache Avro]: https://avro.apache.org
15///
16/// # Example
17///
18/// ```
19/// use polars_core::prelude::*;
20/// use polars_io::avro::AvroWriter;
21/// use std::fs::File;
22/// use polars_io::SerWriter;
23///
24/// fn example(df: &mut DataFrame) -> PolarsResult<()> {
25///     let mut file = File::create("file.avro").expect("could not create file");
26///
27///     AvroWriter::new(&mut file)
28///         .finish(df)
29/// }
30/// ```
31#[must_use]
32pub struct AvroWriter<W> {
33    writer: W,
34    compression: Option<AvroCompression>,
35    name: String,
36}
37
38impl<W> AvroWriter<W>
39where
40    W: Write,
41{
42    /// Set the compression used. Defaults to None.
43    pub fn with_compression(mut self, compression: Option<AvroCompression>) -> Self {
44        self.compression = compression;
45        self
46    }
47
48    pub fn with_name(mut self, name: String) -> Self {
49        self.name = name;
50        self
51    }
52}
53
54impl<W> SerWriter<W> for AvroWriter<W>
55where
56    W: Write,
57{
58    fn new(writer: W) -> Self {
59        Self {
60            writer,
61            compression: None,
62            name: "".to_string(),
63        }
64    }
65
66    fn finish(&mut self, df: &mut DataFrame) -> PolarsResult<()> {
67        let schema = schema_to_arrow_checked(df.schema(), CompatLevel::oldest(), "avro")?;
68        let record = write::to_record(&schema, self.name.clone())?;
69
70        avro_schema::write::write_metadata(&mut self.writer, record.clone(), self.compression)
71            .map_err(to_compute_err)?;
72
73        // Slicing list children to their offsets avoids converting and serializing unused values.
74        let df = DataFrame::new(
75            df.height(),
76            df.apply_columns(|c| {
77                c.trim_lists_to_normalized_offsets()
78                    .unwrap_or_else(|| c.clone())
79            }),
80        )?;
81
82        let mut data = vec![];
83        let mut compressed_block = avro_schema::file::CompressedBlock::default();
84        for chunk in df.iter_chunks(CompatLevel::oldest(), true) {
85            let mut serializers = chunk
86                .iter()
87                .zip(record.fields.iter())
88                .map(|(array, field)| write::new_serializer(array.as_ref(), &field.schema))
89                .collect::<Vec<_>>();
90
91            let mut block =
92                avro_schema::file::Block::new(chunk.arrays()[0].len(), std::mem::take(&mut data));
93            write::serialize(&mut serializers, &mut block);
94            let _was_compressed =
95                avro_schema::write::compress(&mut block, &mut compressed_block, self.compression)
96                    .map_err(to_compute_err)?;
97
98            avro_schema::write::write_block(&mut self.writer, &compressed_block)
99                .map_err(to_compute_err)?;
100            // reuse block for next iteration.
101            data = block.data;
102            data.clear();
103
104            // reuse block for next iteration
105            compressed_block.data.clear();
106            compressed_block.number_of_rows = 0
107        }
108
109        Ok(())
110    }
111}
polars_io/avro/write.rs

polars_io/avro/
write.rs
