query-forge 0.9.0

Run SQL queries and dataset diffs on XLSX/XML/CSV/JSON/JSONL/Markdown/HTML/Feather/Parquet inputs and export results as text, CSV, JSONL, Markdown, XML, HTML, XLSX, Feather, or Parquet
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

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