Martino Ferrari
45 lines
1.3 KiB
Rust
45 lines
1.3 KiB
Rust
use anyhow::Result;
|
|
use parquet::file::properties::WriterProperties;
|
|
use parquet::schema::parser::parse_message_type;
|
|
use parquet::record::RecordWriter;
|
|
use std::fs::File;
|
|
use std::sync::Arc;
|
|
|
|
#[derive(parquet_derive::ParquetRecordWriter)]
|
|
pub struct ScopeRecord {
|
|
pub timestamp: f64,
|
|
pub channel_1: f32,
|
|
}
|
|
|
|
pub struct ParquetWriter {
|
|
writer: parquet::file::writer::SerializedFileWriter<File>,
|
|
}
|
|
|
|
impl ParquetWriter {
|
|
pub fn new(path: &str) -> Result<Self> {
|
|
let file = File::create(path)?;
|
|
let schema = "
|
|
message schema {
|
|
REQUIRED DOUBLE timestamp;
|
|
REQUIRED FLOAT channel_1;
|
|
}
|
|
";
|
|
let schema = Arc::new(parse_message_type(schema)?);
|
|
let props = Arc::new(WriterProperties::builder().build());
|
|
let writer = parquet::file::writer::SerializedFileWriter::new(file, schema, props)?;
|
|
Ok(Self { writer })
|
|
}
|
|
|
|
pub fn write_batch(&mut self, records: &[ScopeRecord]) -> Result<()> {
|
|
let mut row_group_writer = self.writer.next_row_group()?;
|
|
records.write_to_row_group(&mut row_group_writer)?;
|
|
row_group_writer.close()?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn close(self) -> Result<()> {
|
|
self.writer.close()?;
|
|
Ok(())
|
|
}
|
|
}
|