use std::{fs, path::Path, sync::Arc};
use anyhow::{Context, Result};
use arrow_array::{ArrayRef, Float64Array, Int64Array, RecordBatch, StringArray};
use arrow_ipc::writer::FileWriter;
use arrow_schema::{DataType, Field, Schema};
use rust_xlsxwriter::{Format, Workbook};
use crate::{QueryResult, QueryValue};
pub fn write_xlsx(result: &QueryResult, output_path: &Path) -> Result<()> {
const MIN_COLUMN_WIDTH: f64 = 8.0;
const MAX_COLUMN_WIDTH: f64 = 60.0;
const COLUMN_PADDING: f64 = 2.0;
let mut workbook = Workbook::new();
let worksheet = workbook.add_worksheet();
let header_format = Format::new().set_bold();
let mut max_widths: Vec<usize> = result
.columns
.iter()
.map(|header| header.chars().count())
.collect();
for (column, header) in result.columns.iter().enumerate() {
worksheet.write_string_with_format(0, column as u16, header, &header_format)?;
}
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;
if let Some(current_max) = max_widths.get_mut(column_index) {
*current_max = (*current_max).max(query_value_display_width(value));
}
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)?;
}
}
}
}
for (column_index, max_width) in max_widths.iter().enumerate() {
let width = ((*max_width as f64) + COLUMN_PADDING).clamp(MIN_COLUMN_WIDTH, MAX_COLUMN_WIDTH);
worksheet.set_column_width(column_index as u16, width)?;
}
if !result.columns.is_empty() {
let last_column = (result.columns.len() - 1) as u16;
let last_row = result.rows.len() as u32;
worksheet.autofilter(0, 0, last_row, last_column)?;
}
workbook
.save(output_path)
.with_context(|| format!("failed to write {}", output_path.display()))
}
fn query_value_display_width(value: &QueryValue) -> usize {
match value {
QueryValue::Null => 0,
QueryValue::Integer(number) => number.to_string().chars().count(),
QueryValue::Real(number) => number.to_string().chars().count(),
QueryValue::Text(text) => text.chars().count(),
}
}
pub fn write_parquet(result: &QueryResult, output_path: &Path) -> Result<()> {
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(())
}
pub fn write_feather(result: &QueryResult, output_path: &Path) -> Result<()> {
#[derive(Clone, Copy)]
enum ColKind {
Int64,
Double,
Utf8,
}
let col_kinds: Vec<ColKind> = (0..result.columns.len())
.map(|column_index| {
let mut has_real = false;
let mut has_text = false;
for row in &result.rows {
match row.get(column_index).unwrap_or(&QueryValue::Null) {
QueryValue::Real(_) => has_real = true,
QueryValue::Text(_) => has_text = true,
_ => {}
}
}
if has_text {
ColKind::Utf8
} else if has_real {
ColKind::Double
} else {
ColKind::Int64
}
})
.collect();
let schema = Arc::new(Schema::new(
result
.columns
.iter()
.zip(&col_kinds)
.map(|(name, kind)| {
let data_type = match kind {
ColKind::Int64 => DataType::Int64,
ColKind::Double => DataType::Float64,
ColKind::Utf8 => DataType::Utf8,
};
Field::new(name, data_type, true)
})
.collect::<Vec<_>>(),
));
let arrays = col_kinds
.iter()
.enumerate()
.map(|(column_index, kind)| match kind {
ColKind::Int64 => {
let values = result
.rows
.iter()
.map(
|row| match row.get(column_index).unwrap_or(&QueryValue::Null) {
QueryValue::Integer(value) => Some(*value),
QueryValue::Null => None,
QueryValue::Real(_) | QueryValue::Text(_) => None,
},
)
.collect::<Vec<_>>();
Arc::new(Int64Array::from(values)) as ArrayRef
}
ColKind::Double => {
let values = result
.rows
.iter()
.map(
|row| match row.get(column_index).unwrap_or(&QueryValue::Null) {
QueryValue::Integer(value) => Some(*value as f64),
QueryValue::Real(value) => Some(*value),
QueryValue::Null | QueryValue::Text(_) => None,
},
)
.collect::<Vec<_>>();
Arc::new(Float64Array::from(values)) as ArrayRef
}
ColKind::Utf8 => {
let values = result
.rows
.iter()
.map(
|row| match row.get(column_index).unwrap_or(&QueryValue::Null) {
QueryValue::Integer(value) => Some(value.to_string()),
QueryValue::Real(value) => Some(value.to_string()),
QueryValue::Text(value) => Some(value.clone()),
QueryValue::Null => None,
},
)
.collect::<Vec<_>>();
Arc::new(StringArray::from(values)) as ArrayRef
}
})
.collect::<Vec<_>>();
let batch = RecordBatch::try_new(schema.clone(), arrays)
.context("failed to build Feather record batch")?;
let file = fs::File::create(output_path)
.with_context(|| format!("failed to create {}", output_path.display()))?;
let mut writer =
FileWriter::try_new(file, &schema).context("failed to initialize Feather writer")?;
writer
.write(&batch)
.context("failed to write Feather record batch")?;
writer
.finish()
.context("failed to finalize Feather writer")?;
Ok(())
}