use std::fs;
use std::io::BufWriter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use crate::{Error, Job, Result};
pub fn parquet_to_csv(job: &Job) -> Result<()> {
parquet_to_sv(job, b',')
}
pub fn parquet_to_tsv(job: &Job) -> Result<()> {
parquet_to_sv(job, b'\t')
}
pub fn parquet_to_json(job: &Job) -> Result<()> {
let reader = open_reader(job)?;
let file = fs::File::create(job.output)?;
let mut writer = arrow_json::ArrayWriter::new(BufWriter::new(file));
for batch in reader {
let batch =
batch.map_err(|e| Error::other(format!("parquet read: {e}")))?;
writer
.write(&batch)
.map_err(|e| Error::other(format!("json write: {e}")))?;
}
writer
.finish()
.map_err(|e| Error::other(format!("json finish: {e}")))?;
Ok(())
}
fn parquet_to_sv(job: &Job, delimiter: u8) -> Result<()> {
let reader = open_reader(job)?;
let file = fs::File::create(job.output)?;
let mut writer = arrow_csv::WriterBuilder::new()
.with_header(true)
.with_delimiter(delimiter)
.build(BufWriter::new(file));
for batch in reader {
let batch =
batch.map_err(|e| Error::other(format!("parquet read: {e}")))?;
writer
.write(&batch)
.map_err(|e| Error::other(format!("csv write: {e}")))?;
}
Ok(())
}
fn open_reader(
job: &Job,
) -> Result<parquet::arrow::arrow_reader::ParquetRecordBatchReader> {
let file = fs::File::open(job.input)?;
let builder = ParquetRecordBatchReaderBuilder::try_new(file)
.map_err(|e| Error::other(format!("parquet open: {e}")))?;
builder
.build()
.map_err(|e| Error::other(format!("parquet reader: {e}")))
}