use std::{fs, path::Path};
use anyhow::{Context, Result};
use rust_xlsxwriter::Workbook;
use crate::{QueryResult, QueryValue};
pub fn write_xlsx(result: &QueryResult, output_path: &Path) -> Result<()> {
let mut workbook = Workbook::new();
let worksheet = workbook.add_worksheet();
for (column, header) in result.columns.iter().enumerate() {
worksheet.write_string(0, column as u16, header)?;
}
for (row_index, row) in result.rows.iter().enumerate() {
for (column_index, value) in row.iter().enumerate() {
let excel_row = (row_index + 1) as u32;
let excel_column = column_index as u16;
match value {
QueryValue::Null => {}
QueryValue::Integer(number) => {
worksheet.write_number(excel_row, excel_column, *number as f64)?;
}
QueryValue::Real(number) => {
worksheet.write_number(excel_row, excel_column, *number)?;
}
QueryValue::Text(text) => {
worksheet.write_string(excel_row, excel_column, text)?;
}
}
}
}
workbook
.save(output_path)
.with_context(|| format!("failed to write {}", output_path.display()))
}
pub fn write_parquet(result: &QueryResult, output_path: &Path) -> Result<()> {
use std::sync::Arc;
use parquet::basic::{ConvertedType, Repetition, Type as PhysicalType};
use parquet::column::writer::ColumnWriter;
use parquet::data_type::ByteArray;
use parquet::file::properties::WriterProperties;
use parquet::file::writer::SerializedFileWriter;
use parquet::schema::types::Type as ParquetType;
#[derive(Clone, Copy)]
enum ColKind {
Int64,
Double,
Bytes,
}
let col_kinds: Vec<ColKind> = (0..result.columns.len())
.map(|ci| {
let mut has_real = false;
let mut has_text = false;
for row in &result.rows {
match row.get(ci).unwrap_or(&QueryValue::Null) {
QueryValue::Real(_) => has_real = true,
QueryValue::Text(_) => has_text = true,
_ => {}
}
}
if has_text {
ColKind::Bytes
} else if has_real {
ColKind::Double
} else {
ColKind::Int64
}
})
.collect();
let fields: Vec<Arc<ParquetType>> = result
.columns
.iter()
.zip(&col_kinds)
.map(|(name, kind)| {
let physical = match kind {
ColKind::Int64 => PhysicalType::INT64,
ColKind::Double => PhysicalType::DOUBLE,
ColKind::Bytes => PhysicalType::BYTE_ARRAY,
};
let mut builder = ParquetType::primitive_type_builder(name, physical)
.with_repetition(Repetition::OPTIONAL);
if matches!(kind, ColKind::Bytes) {
builder = builder.with_converted_type(ConvertedType::UTF8);
}
Arc::new(
builder
.build()
.with_context(|| format!("failed to build Parquet field '{name}'"))
.expect("valid field"),
)
})
.collect();
let schema = Arc::new(
ParquetType::group_type_builder("schema")
.with_fields(fields)
.build()
.context("failed to build Parquet schema")?,
);
let props = Arc::new(WriterProperties::builder().build());
let file = fs::File::create(output_path)
.with_context(|| format!("failed to create {}", output_path.display()))?;
let mut file_writer = SerializedFileWriter::new(file, schema, props)
.context("failed to initialize Parquet writer")?;
let mut rg = file_writer
.next_row_group()
.context("failed to start Parquet row group")?;
for (col_idx, kind) in col_kinds.iter().enumerate() {
let def_levels: Vec<i16> = result
.rows
.iter()
.map(|row| match row.get(col_idx).unwrap_or(&QueryValue::Null) {
QueryValue::Null => 0,
_ => 1,
})
.collect();
let Some(mut col_writer) = rg
.next_column()
.context("failed to open Parquet column writer")?
else {
break;
};
match (kind, col_writer.untyped()) {
(ColKind::Int64, ColumnWriter::Int64ColumnWriter(w)) => {
let values: Vec<i64> = result
.rows
.iter()
.filter_map(|row| match row.get(col_idx).unwrap_or(&QueryValue::Null) {
QueryValue::Integer(v) => Some(*v),
_ => None,
})
.collect();
w.write_batch(&values, Some(&def_levels), None)?;
}
(ColKind::Double, ColumnWriter::DoubleColumnWriter(w)) => {
let values: Vec<f64> = result
.rows
.iter()
.filter_map(|row| match row.get(col_idx).unwrap_or(&QueryValue::Null) {
QueryValue::Real(v) => Some(*v),
QueryValue::Integer(v) => Some(*v as f64),
_ => None,
})
.collect();
w.write_batch(&values, Some(&def_levels), None)?;
}
(ColKind::Bytes, ColumnWriter::ByteArrayColumnWriter(w)) => {
let values: Vec<ByteArray> = result
.rows
.iter()
.filter_map(|row| match row.get(col_idx).unwrap_or(&QueryValue::Null) {
QueryValue::Text(s) => Some(ByteArray::from(s.as_bytes().to_vec())),
QueryValue::Integer(v) => Some(ByteArray::from(v.to_string().into_bytes())),
QueryValue::Real(v) => Some(ByteArray::from(v.to_string().into_bytes())),
_ => None,
})
.collect();
w.write_batch(&values, Some(&def_levels), None)?;
}
_ => unreachable!("column kind and writer type must agree"),
}
col_writer.close()?;
}
rg.close()?;
file_writer.close()?;
Ok(())
}