use arrow2::{
array::Array,
chunk::Chunk,
datatypes::Schema,
io::parquet::write::{
transverse, CompressionOptions, Encoding, FileWriter, RowGroupIterator, Version,
WriteOptions,
},
};
use std::fs::File;
pub(crate) fn write_chunk_to_file(
path: &str,
schema: Schema,
chunk: Chunk<Box<dyn Array>>,
) -> anyhow::Result<()> {
let options = WriteOptions {
write_statistics: true,
compression: CompressionOptions::Zstd(None),
version: Version::V2,
data_pagesize_limit: None,
};
let file = File::create(path)?;
let mut writer = FileWriter::try_new(file, schema.clone(), options)?;
let iter = vec![Ok(chunk)];
let encodings = schema
.fields
.iter()
.map(|f| transverse(&f.data_type, |_| Encoding::Plain))
.collect::<Vec<Vec<Encoding>>>();
let row_groups = RowGroupIterator::try_new(iter.into_iter(), &schema, options, encodings)?;
for group in row_groups {
writer.write(group?)?;
}
let _size = writer.end(None)?;
Ok(())
}