query-forge 0.9.0

use anyhow::{Context, Result, bail};
use arrow_array::{
    Array, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array,
    LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, UInt8Array, UInt16Array,
    UInt32Array, UInt64Array,
};
use arrow_cast::display::array_value_to_string;
use arrow_ipc::reader::FileReader;
use arrow_schema::DataType;

use crate::{QueryValue, TypeInferenceOptions};

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

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

    let file = std::fs::File::open(feather_path)
        .with_context(|| format!("failed to open {}", feather_path.display()))?;
    let reader = FileReader::try_new(file, None)
        .with_context(|| format!("failed to read Feather file {}", feather_path.display()))?;

    let columns = normalize_text_headers(
        &reader
            .schema()
            .fields()
            .iter()
            .map(|field| field.name().clone())
            .collect::<Vec<_>>(),
    );

    let mut rows = Vec::new();
    for batch in reader {
        let batch = batch
            .with_context(|| format!("failed to read batch from {}", feather_path.display()))?;
        rows.extend(record_batch_to_rows(&batch, inference_options)?);
    }

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

fn record_batch_to_rows(
    batch: &RecordBatch,
    inference_options: &TypeInferenceOptions,
) -> Result<Vec<Vec<QueryValue>>> {
    let mut rows = Vec::with_capacity(batch.num_rows());

    for row_index in 0..batch.num_rows() {
        let mut row = Vec::with_capacity(batch.num_columns());
        for column in batch.columns() {
            row.push(arrow_value_to_query_value(
                column.as_ref(),
                row_index,
                inference_options,
            )?);
        }
        rows.push(row);
    }

    Ok(rows)
}

fn arrow_value_to_query_value(
    array: &dyn Array,
    row_index: usize,
    inference_options: &TypeInferenceOptions,
) -> Result<QueryValue> {
    if array.is_null(row_index) {
        return Ok(QueryValue::Null);
    }

    let value = match array.data_type() {
        DataType::Boolean => QueryValue::Integer(i64::from(
            array
                .as_any()
                .downcast_ref::<BooleanArray>()
                .expect("boolean array")
                .value(row_index),
        )),
        DataType::Int8 => QueryValue::Integer(i64::from(
            array
                .as_any()
                .downcast_ref::<Int8Array>()
                .expect("int8 array")
                .value(row_index),
        )),
        DataType::Int16 => QueryValue::Integer(i64::from(
            array
                .as_any()
                .downcast_ref::<Int16Array>()
                .expect("int16 array")
                .value(row_index),
        )),
        DataType::Int32 => QueryValue::Integer(i64::from(
            array
                .as_any()
                .downcast_ref::<Int32Array>()
                .expect("int32 array")
                .value(row_index),
        )),
        DataType::Int64 => QueryValue::Integer(
            array
                .as_any()
                .downcast_ref::<Int64Array>()
                .expect("int64 array")
                .value(row_index),
        ),
        DataType::UInt8 => QueryValue::Integer(i64::from(
            array
                .as_any()
                .downcast_ref::<UInt8Array>()
                .expect("uint8 array")
                .value(row_index),
        )),
        DataType::UInt16 => QueryValue::Integer(i64::from(
            array
                .as_any()
                .downcast_ref::<UInt16Array>()
                .expect("uint16 array")
                .value(row_index),
        )),
        DataType::UInt32 => QueryValue::Integer(i64::from(
            array
                .as_any()
                .downcast_ref::<UInt32Array>()
                .expect("uint32 array")
                .value(row_index),
        )),
        DataType::UInt64 => {
            let value = array
                .as_any()
                .downcast_ref::<UInt64Array>()
                .expect("uint64 array")
                .value(row_index);
            i64::try_from(value)
                .map(QueryValue::Integer)
                .unwrap_or_else(|_| QueryValue::Text(value.to_string()))
        }
        DataType::Float32 => QueryValue::Real(f64::from(
            array
                .as_any()
                .downcast_ref::<Float32Array>()
                .expect("float32 array")
                .value(row_index),
        )),
        DataType::Float64 => QueryValue::Real(
            array
                .as_any()
                .downcast_ref::<Float64Array>()
                .expect("float64 array")
                .value(row_index),
        ),
        DataType::Utf8 => QueryValue::Text(
            array
                .as_any()
                .downcast_ref::<StringArray>()
                .expect("utf8 array")
                .value(row_index)
                .to_owned(),
        ),
        DataType::LargeUtf8 => QueryValue::Text(
            array
                .as_any()
                .downcast_ref::<LargeStringArray>()
                .expect("large utf8 array")
                .value(row_index)
                .to_owned(),
        ),
        DataType::Binary => QueryValue::Text(
            String::from_utf8_lossy(
                array
                    .as_any()
                    .downcast_ref::<arrow_array::BinaryArray>()
                    .expect("binary array")
                    .value(row_index),
            )
            .into_owned(),
        ),
        DataType::LargeBinary => QueryValue::Text(
            String::from_utf8_lossy(
                array
                    .as_any()
                    .downcast_ref::<LargeBinaryArray>()
                    .expect("large binary array")
                    .value(row_index),
            )
            .into_owned(),
        ),
        _ => QueryValue::Text(
            array_value_to_string(array, row_index)
                .with_context(|| format!("failed to render Feather value at row {row_index}"))?,
        ),
    };

    Ok(apply_inference_overrides(value, inference_options))
}