mongreldb_kit/
arrow_util.rs

1//! Arrow IPC <-> JSON row helpers, shared by the embedded SQL surface
2//! ([`Database::sql`] / `sql_arrow` / `sql_rows`) and the optional `remote`
3//! HTTP client. Kept here (rather than inside `remote.rs`) so the default
4//! embedded build — which has no HTTP dependency — can still decode the Arrow
5//! IPC bytes that `MongrelSession::run` produces.
6
7use arrow::array::{
8    Array, BooleanArray, Float64Array, Int32Array, Int64Array, NullArray, StringArray,
9};
10use arrow::ipc::reader::FileReader;
11use arrow::ipc::writer::FileWriter;
12use arrow::record_batch::RecordBatch;
13use serde_json::{json, Map, Value};
14
15use crate::error::{KitError, Result};
16
17/// Decode Arrow IPC *file* bytes (the format `MongrelSession` and the daemon
18/// both emit) into [`RecordBatch`]es. An empty input yields an empty vec.
19pub fn read_arrow_ipc(bytes: &[u8]) -> Result<Vec<RecordBatch>> {
20    if bytes.is_empty() {
21        return Ok(Vec::new());
22    }
23    let cursor = std::io::Cursor::new(bytes);
24    let reader = FileReader::try_new(cursor, None)
25        .map_err(|e| KitError::Storage(format!("arrow ipc decode: {e}")))?;
26    reader
27        .into_iter()
28        .collect::<std::result::Result<Vec<_>, _>>()
29        .map_err(|e| KitError::Storage(format!("arrow ipc decode: {e}")))
30}
31
32/// Encode `RecordBatch`es as Arrow IPC *file* bytes — the wire format the
33/// daemon and the NAPI addon both produce. Mirrors the node addon's
34/// `native_cols_to_ipc_from_batches`.
35pub fn batches_to_ipc(batches: &[RecordBatch]) -> Result<Vec<u8>> {
36    let schema = batches
37        .first()
38        .map(|b| b.schema())
39        .unwrap_or_else(|| std::sync::Arc::new(arrow::datatypes::Schema::empty()));
40    let mut out = Vec::new();
41    let mut writer =
42        FileWriter::try_new(&mut out, &schema).map_err(|e| KitError::Storage(e.to_string()))?;
43    for batch in batches {
44        writer
45            .write(batch)
46            .map_err(|e| KitError::Storage(format!("arrow ipc encode: {e}")))?;
47    }
48    writer
49        .finish()
50        .map_err(|e| KitError::Storage(format!("arrow ipc encode: {e}")))?;
51    Ok(out)
52}
53
54/// Materialize one `RecordBatch` into JSON-row maps (column name → value).
55pub fn batch_to_rows(b: &RecordBatch) -> Result<Vec<Map<String, Value>>> {
56    let schema = b.schema();
57    let mut rows = Vec::with_capacity(b.num_rows());
58    for r in 0..b.num_rows() {
59        let mut row = Map::new();
60        for (c, field) in schema.fields().iter().enumerate() {
61            let name = field.name();
62            let arr = b.column(c);
63            row.insert(name.clone(), cell_value(arr.as_ref(), r));
64        }
65        rows.push(row);
66    }
67    Ok(rows)
68}
69
70/// Flatten a slice of batches into a single list of JSON-row maps.
71pub fn batches_to_rows(batches: &[RecordBatch]) -> Result<Vec<Map<String, Value>>> {
72    let mut rows = Vec::new();
73    for batch in batches {
74        rows.extend(batch_to_rows(batch)?);
75    }
76    Ok(rows)
77}
78
79fn cell_value(arr: &dyn Array, r: usize) -> Value {
80    if arr.is_null(r) {
81        return Value::Null;
82    }
83    if let Some(a) = arr.as_any().downcast_ref::<Int64Array>() {
84        return json!(a.value(r));
85    }
86    if let Some(a) = arr.as_any().downcast_ref::<Int32Array>() {
87        return json!(a.value(r));
88    }
89    if let Some(a) = arr.as_any().downcast_ref::<Float64Array>() {
90        return serde_json::Number::from_f64(a.value(r))
91            .map(Value::Number)
92            .unwrap_or(Value::Null);
93    }
94    if let Some(a) = arr.as_any().downcast_ref::<BooleanArray>() {
95        return Value::Bool(a.value(r));
96    }
97    if let Some(a) = arr.as_any().downcast_ref::<StringArray>() {
98        return Value::String(a.value(r).to_string());
99    }
100    if arr.as_any().downcast_ref::<NullArray>().is_some() {
101        return Value::Null;
102    }
103    // Fallback: stringify unknown types.
104    Value::String(format!("{arr:?}"))
105}
mongreldb_kit/arrow_util.rs

mongreldb_kit/
arrow_util.rs
