query_forge/

writers.rs

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,
    }

    // Infer the physical type for each column from its values.
    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();

    // Build the Parquet schema.
    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(())
}