sheeter-engine 1.5.0

//! Parquet reading and writing.
//!
//! Reading uses the `parquet` crate's Arrow bridge. Each row becomes a
//! [`CanonicalValue::Object`] whose keys are the Parquet column names.
//! Supported Arrow types: signed/unsigned 8/16/32/64-bit integers, 32/64-bit
//! floats, booleans, UTF-8, and large UTF-8. Anything else is stringified via
//! its `Debug` representation, so an exotic schema degrades gracefully
//! instead of failing the whole read.
//!
//! Writing infers an Arrow schema from the first batch of rows. Per column:
//! all-integer numbers become Int64, all-float numbers become Float64,
//! boolean-only values become Boolean, and every other shape (mixed types,
//! strings, nested values) becomes Utf8. Unsigned 64-bit values larger than
//! `i64::MAX` fall back to Utf8 to avoid silent truncation. Snappy
//! compression is used by default.
//!
//! ```ignore
//! // Round-trip covered in tests/parquet.rs.
//! ```

use crate::CanonicalValue;
use std::error::Error;
use std::sync::Arc;

/// Parses a Parquet byte buffer into an array of row objects.
///
/// # Errors
///
/// Returns an error if the buffer is not valid Parquet, if the Arrow reader
/// rejects the schema, or if any batch cannot be traversed.
pub fn from_parquet(bytes: &[u8]) -> Result<CanonicalValue, Box<dyn Error>> {
    use arrow::array::Array;
    use arrow::record_batch::RecordBatch;
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

    let data = bytes::Bytes::copy_from_slice(bytes);
    let builder = ParquetRecordBatchReaderBuilder::try_new(data)?;
    let schema = builder.schema().clone();
    let reader = builder.build()?;

    let mut rows: Vec<CanonicalValue> = Vec::new();
    for batch_result in reader {
        let batch: RecordBatch = batch_result?;
        for r in 0..batch.num_rows() {
            let mut obj: Vec<(String, CanonicalValue)> = Vec::with_capacity(schema.fields().len());
            for c in 0..schema.fields().len() {
                let col = batch.column(c);
                let name = schema.field(c).name().clone();
                obj.push((name, arrow_value(col.as_ref(), r)));
            }
            rows.push(CanonicalValue::Object(obj));
        }
    }
    Ok(CanonicalValue::Array(rows))
}

fn arrow_value(col: &dyn arrow::array::Array, r: usize) -> CanonicalValue {
    use arrow::array::{
        BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array,
        LargeStringArray, StringArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
    };
    if col.is_null(r) {
        return CanonicalValue::Null;
    }
    if let Some(a) = col.as_any().downcast_ref::<Int8Array>() {
        return CanonicalValue::Number(serde_json::Number::from(a.value(r) as i64));
    }
    if let Some(a) = col.as_any().downcast_ref::<Int16Array>() {
        return CanonicalValue::Number(serde_json::Number::from(a.value(r) as i64));
    }
    if let Some(a) = col.as_any().downcast_ref::<Int32Array>() {
        return CanonicalValue::Number(serde_json::Number::from(a.value(r) as i64));
    }
    if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
        return CanonicalValue::Number(serde_json::Number::from(a.value(r)));
    }
    if let Some(a) = col.as_any().downcast_ref::<UInt8Array>() {
        return CanonicalValue::Number(serde_json::Number::from(a.value(r) as u64));
    }
    if let Some(a) = col.as_any().downcast_ref::<UInt16Array>() {
        return CanonicalValue::Number(serde_json::Number::from(a.value(r) as u64));
    }
    if let Some(a) = col.as_any().downcast_ref::<UInt32Array>() {
        return CanonicalValue::Number(serde_json::Number::from(a.value(r) as u64));
    }
    if let Some(a) = col.as_any().downcast_ref::<UInt64Array>() {
        return CanonicalValue::Number(serde_json::Number::from(a.value(r)));
    }
    if let Some(a) = col.as_any().downcast_ref::<Float32Array>() {
        return serde_json::Number::from_f64(a.value(r) as f64)
            .map(CanonicalValue::Number)
            .unwrap_or(CanonicalValue::Null);
    }
    if let Some(a) = col.as_any().downcast_ref::<Float64Array>() {
        return serde_json::Number::from_f64(a.value(r))
            .map(CanonicalValue::Number)
            .unwrap_or(CanonicalValue::Null);
    }
    if let Some(a) = col.as_any().downcast_ref::<BooleanArray>() {
        return CanonicalValue::Bool(a.value(r));
    }
    if let Some(a) = col.as_any().downcast_ref::<StringArray>() {
        return CanonicalValue::String(a.value(r).to_string());
    }
    if let Some(a) = col.as_any().downcast_ref::<LargeStringArray>() {
        return CanonicalValue::String(a.value(r).to_string());
    }
    CanonicalValue::String(format!("{:?}", col.slice(r, 1)))
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ColumnType {
    Int64,
    Float64,
    Boolean,
    Utf8,
}

fn infer_column(values: &[Option<&CanonicalValue>]) -> ColumnType {
    let mut has_int = false;
    let mut has_float = false;
    let mut has_bool = false;
    let mut has_string = false;
    let mut u64_overflow = false;
    for v in values.iter().flatten() {
        match v {
            CanonicalValue::Number(n) => {
                if n.is_f64() {
                    has_float = true;
                } else {
                    has_int = true;
                    if let Some(u) = n.as_u64() {
                        if u > i64::MAX as u64 {
                            u64_overflow = true;
                        }
                    }
                }
            }
            CanonicalValue::Bool(_) => has_bool = true,
            CanonicalValue::String(_) => has_string = true,
            CanonicalValue::Null => {}
            _ => has_string = true,
        }
    }
    if has_string {
        return ColumnType::Utf8;
    }
    if has_int && has_bool {
        return ColumnType::Utf8;
    }
    if has_float && has_bool {
        return ColumnType::Utf8;
    }
    if has_float {
        return ColumnType::Float64;
    }
    if has_int {
        return if u64_overflow {
            ColumnType::Utf8
        } else {
            ColumnType::Int64
        };
    }
    if has_bool {
        return ColumnType::Boolean;
    }
    ColumnType::Utf8
}

fn arrow_datatype(t: ColumnType) -> arrow::datatypes::DataType {
    use arrow::datatypes::DataType;
    match t {
        ColumnType::Int64 => DataType::Int64,
        ColumnType::Float64 => DataType::Float64,
        ColumnType::Boolean => DataType::Boolean,
        ColumnType::Utf8 => DataType::Utf8,
    }
}

fn build_array(
    col_type: ColumnType,
    values: &[Option<&CanonicalValue>],
) -> Result<arrow::array::ArrayRef, Box<dyn Error>> {
    use arrow::array::{ArrayRef, BooleanBuilder, Float64Builder, Int64Builder, StringBuilder};
    match col_type {
        ColumnType::Int64 => {
            let mut b = Int64Builder::with_capacity(values.len());
            for v in values {
                match v {
                    Some(CanonicalValue::Number(n)) => {
                        if let Some(i) = n.as_i64() {
                            b.append_value(i);
                        } else if let Some(u) = n.as_u64() {
                            b.append_value(u as i64);
                        } else {
                            b.append_null();
                        }
                    }
                    _ => b.append_null(),
                }
            }
            Ok(Arc::new(b.finish()) as ArrayRef)
        }
        ColumnType::Float64 => {
            let mut b = Float64Builder::with_capacity(values.len());
            for v in values {
                match v {
                    Some(CanonicalValue::Number(n)) => match n.as_f64() {
                        Some(f) => b.append_value(f),
                        None => b.append_null(),
                    },
                    _ => b.append_null(),
                }
            }
            Ok(Arc::new(b.finish()) as ArrayRef)
        }
        ColumnType::Boolean => {
            let mut b = BooleanBuilder::with_capacity(values.len());
            for v in values {
                match v {
                    Some(CanonicalValue::Bool(x)) => b.append_value(*x),
                    _ => b.append_null(),
                }
            }
            Ok(Arc::new(b.finish()) as ArrayRef)
        }
        ColumnType::Utf8 => {
            let mut b = StringBuilder::with_capacity(values.len(), values.len() * 8);
            for v in values {
                match v {
                    Some(CanonicalValue::Null) | None => b.append_null(),
                    Some(CanonicalValue::String(s)) => b.append_value(s.as_str()),
                    Some(CanonicalValue::Bool(x)) => b.append_value(x.to_string()),
                    Some(CanonicalValue::Number(n)) => b.append_value(n.to_string()),
                    Some(other) => b.append_value(format!("{other:?}")),
                }
            }
            Ok(Arc::new(b.finish()) as ArrayRef)
        }
    }
}

/// Serializes a `CanonicalValue::Array<Object>` to Parquet.
///
/// # Errors
///
/// Returns an error if the value is not tabular (neither an array of objects
/// nor a single object), if the array is empty (there is nothing to infer a
/// schema from), or if the Arrow writer fails.
pub fn to_parquet(value: &CanonicalValue) -> Result<Vec<u8>, Box<dyn Error>> {
    use arrow::datatypes::{Field, Schema};
    use arrow::record_batch::RecordBatch;
    use parquet::arrow::ArrowWriter;
    use parquet::basic::Compression;
    use parquet::file::properties::WriterProperties;

    let rows: Vec<&Vec<(String, CanonicalValue)>> = match value {
        CanonicalValue::Array(a) => a
            .iter()
            .map(|r| match r {
                CanonicalValue::Object(o) => Ok(o),
                _ => Err("Parquet requires an array of objects".to_string()),
            })
            .collect::<Result<Vec<_>, _>>()?,
        CanonicalValue::Object(o) => vec![o],
        _ => return Err("Parquet requires tabular data".into()),
    };
    if rows.is_empty() {
        return Err("cannot write empty Parquet (no rows to infer schema from)".into());
    }

    let mut headers: Vec<String> = Vec::new();
    for row in &rows {
        for (k, _) in row.iter() {
            if !headers.contains(k) {
                headers.push(k.clone());
            }
        }
    }

    let mut column_values: Vec<Vec<Option<&CanonicalValue>>> = Vec::with_capacity(headers.len());
    let mut column_types: Vec<ColumnType> = Vec::with_capacity(headers.len());
    for h in &headers {
        let vals: Vec<Option<&CanonicalValue>> = rows
            .iter()
            .map(|row| row.iter().find(|(k, _)| k == h).map(|(_, v)| v))
            .collect();
        column_types.push(infer_column(&vals));
        column_values.push(vals);
    }

    let fields: Vec<Field> = headers
        .iter()
        .zip(column_types.iter())
        .map(|(h, t)| Field::new(h.as_str(), arrow_datatype(*t), true))
        .collect();
    let schema = Arc::new(Schema::new(fields));

    let mut arrays: Vec<arrow::array::ArrayRef> = Vec::with_capacity(headers.len());
    for (i, t) in column_types.iter().enumerate() {
        arrays.push(build_array(*t, &column_values[i])?);
    }

    let batch = RecordBatch::try_new(schema.clone(), arrays)?;

    let props = WriterProperties::builder()
        .set_compression(Compression::SNAPPY)
        .build();

    let mut buf: Vec<u8> = Vec::new();
    {
        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props))?;
        writer.write(&batch)?;
        writer.close()?;
    }
    Ok(buf)
}