query-forge 0.7.0

use anyhow::{Context, Result, bail};
use parquet::{
    file::reader::{FileReader, SerializedFileReader},
    record::Field,
};

use crate::{QueryValue, TypeInferenceOptions};

use super::{SheetData, apply_inference_overrides, normalize_text_headers};

pub(super) fn load_parquet_sheet(
    parquet_path: &std::path::Path,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
) -> Result<SheetData> {
    if let Some(selector) = requested_sheet {
        bail!(
            "Parquet input {} does not support selector '{selector}'. Remove ':{selector}' from --input for Parquet files.",
            parquet_path.display()
        );
    }

    let reader = SerializedFileReader::try_from(parquet_path)
        .with_context(|| format!("failed to open {}", parquet_path.display()))?;
    let mut row_iter = reader
        .get_row_iter(None)
        .with_context(|| format!("failed to read rows from {}", parquet_path.display()))?;

    let first_row = row_iter
        .next()
        .transpose()
        .with_context(|| format!("failed to read first row from {}", parquet_path.display()))?;

    let Some(first_row) = first_row else {
        bail!("Parquet input {} is empty", parquet_path.display());
    };

    let columns = normalize_text_headers(
        &first_row
            .get_column_iter()
            .map(|(name, _)| name.to_owned())
            .collect::<Vec<_>>(),
    );

    let mut rows = vec![parquet_row_to_values(first_row, inference_options)];
    for row in row_iter {
        let row =
            row.with_context(|| format!("failed to read row from {}", parquet_path.display()))?;
        rows.push(parquet_row_to_values(row, inference_options));
    }

    Ok(SheetData {
        original_name: "parquet".to_owned(),
        columns,
        rows,
    })
}

fn parquet_row_to_values(
    row: parquet::record::Row,
    inference_options: &TypeInferenceOptions,
) -> Vec<QueryValue> {
    row.into_columns()
        .into_iter()
        .map(|(_, field)| parquet_field_to_query_value(field, inference_options))
        .collect()
}

fn parquet_field_to_query_value(
    field: Field,
    inference_options: &TypeInferenceOptions,
) -> QueryValue {
    apply_inference_overrides(
        match field {
            Field::Null => QueryValue::Null,
            Field::Bool(value) => QueryValue::Integer(i64::from(value)),
            Field::Byte(value) => QueryValue::Integer(i64::from(value)),
            Field::Short(value) => QueryValue::Integer(i64::from(value)),
            Field::Int(value) => QueryValue::Integer(i64::from(value)),
            Field::Long(value) => QueryValue::Integer(value),
            Field::UByte(value) => QueryValue::Integer(i64::from(value)),
            Field::UShort(value) => QueryValue::Integer(i64::from(value)),
            Field::UInt(value) => QueryValue::Integer(i64::from(value)),
            Field::ULong(value) => i64::try_from(value)
                .map(QueryValue::Integer)
                .unwrap_or_else(|_| QueryValue::Text(value.to_string())),
            Field::Float16(value) => QueryValue::Real(f64::from(value)),
            Field::Float(value) => QueryValue::Real(f64::from(value)),
            Field::Double(value) => QueryValue::Real(value),
            Field::Decimal(value) => QueryValue::Text(format!("{value:?}")),
            Field::Str(value) => QueryValue::Text(value),
            Field::Bytes(value) => QueryValue::Text(String::from_utf8_lossy(value.data()).into()),
            Field::Date(value) => QueryValue::Integer(i64::from(value)),
            Field::TimestampMillis(value) => QueryValue::Integer(value),
            Field::TimestampMicros(value) => QueryValue::Integer(value),
            Field::Group(value) => QueryValue::Text(value.to_string()),
            Field::ListInternal(value) => QueryValue::Text(format!("{value:?}")),
            Field::MapInternal(value) => QueryValue::Text(format!("{value:?}")),
        },
        inference_options,
    )
}