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, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array,
9    NullArray, StringArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
10};
11use arrow::ipc::reader::FileReader;
12use arrow::ipc::writer::FileWriter;
13use arrow::record_batch::RecordBatch;
14use serde_json::{json, Map, Value};
15
16use crate::error::{KitError, Result};
17
18/// Decode Arrow IPC *file* bytes (the format `MongrelSession` and the daemon
19/// both emit) into [`RecordBatch`]es. An empty input yields an empty vec.
20pub fn read_arrow_ipc(bytes: &[u8]) -> Result<Vec<RecordBatch>> {
21    if bytes.is_empty() {
22        return Ok(Vec::new());
23    }
24    let cursor = std::io::Cursor::new(bytes);
25    let reader = FileReader::try_new(cursor, None)
26        .map_err(|e| KitError::Storage(format!("arrow ipc decode: {e}")))?;
27    reader
28        .into_iter()
29        .collect::<std::result::Result<Vec<_>, _>>()
30        .map_err(|e| KitError::Storage(format!("arrow ipc decode: {e}")))
31}
32
33/// Encode `RecordBatch`es as Arrow IPC *file* bytes — the wire format the
34/// daemon and the NAPI addon both produce. Mirrors the node addon's
35/// `native_cols_to_ipc_from_batches`.
36pub fn batches_to_ipc(batches: &[RecordBatch]) -> Result<Vec<u8>> {
37    let schema = batches
38        .first()
39        .map(|b| b.schema())
40        .unwrap_or_else(|| std::sync::Arc::new(arrow::datatypes::Schema::empty()));
41    let mut out = Vec::new();
42    let mut writer =
43        FileWriter::try_new(&mut out, &schema).map_err(|e| KitError::Storage(e.to_string()))?;
44    for batch in batches {
45        writer
46            .write(batch)
47            .map_err(|e| KitError::Storage(format!("arrow ipc encode: {e}")))?;
48    }
49    writer
50        .finish()
51        .map_err(|e| KitError::Storage(format!("arrow ipc encode: {e}")))?;
52    Ok(out)
53}
54
55/// Materialize one `RecordBatch` into JSON-row maps (column name → value).
56pub fn batch_to_rows(b: &RecordBatch) -> Result<Vec<Map<String, Value>>> {
57    let schema = b.schema();
58    let mut rows = Vec::with_capacity(b.num_rows());
59    for r in 0..b.num_rows() {
60        let mut row = Map::new();
61        for (c, field) in schema.fields().iter().enumerate() {
62            let name = field.name();
63            let arr = b.column(c);
64            row.insert(name.clone(), cell_value(arr.as_ref(), r));
65        }
66        rows.push(row);
67    }
68    Ok(rows)
69}
70
71/// Flatten a slice of batches into a single list of JSON-row maps.
72pub fn batches_to_rows(batches: &[RecordBatch]) -> Result<Vec<Map<String, Value>>> {
73    let mut rows = Vec::new();
74    for batch in batches {
75        rows.extend(batch_to_rows(batch)?);
76    }
77    Ok(rows)
78}
79
80fn cell_value(arr: &dyn Array, r: usize) -> Value {
81    if arr.is_null(r) {
82        return Value::Null;
83    }
84    // Signed integers → i64
85    if let Some(a) = arr.as_any().downcast_ref::<Int64Array>() {
86        return json!(a.value(r));
87    }
88    if let Some(a) = arr.as_any().downcast_ref::<Int32Array>() {
89        return json!(a.value(r));
90    }
91    if let Some(a) = arr.as_any().downcast_ref::<Int16Array>() {
92        return json!(a.value(r));
93    }
94    if let Some(a) = arr.as_any().downcast_ref::<Int8Array>() {
95        return json!(a.value(r));
96    }
97    // Unsigned integers → i64 (JSON has no unsigned; cast is lossless for normal values)
98    if let Some(a) = arr.as_any().downcast_ref::<UInt64Array>() {
99        return json!(a.value(r) as i64);
100    }
101    if let Some(a) = arr.as_any().downcast_ref::<UInt32Array>() {
102        return json!(a.value(r) as i64);
103    }
104    if let Some(a) = arr.as_any().downcast_ref::<UInt16Array>() {
105        return json!(a.value(r) as i64);
106    }
107    if let Some(a) = arr.as_any().downcast_ref::<UInt8Array>() {
108        return json!(a.value(r) as i64);
109    }
110    // Floats
111    if let Some(a) = arr.as_any().downcast_ref::<Float64Array>() {
112        return serde_json::Number::from_f64(a.value(r))
113            .map(Value::Number)
114            .unwrap_or(Value::Null);
115    }
116    if let Some(a) = arr.as_any().downcast_ref::<Float32Array>() {
117        return serde_json::Number::from_f64(a.value(r) as f64)
118            .map(Value::Number)
119            .unwrap_or(Value::Null);
120    }
121    // Boolean
122    if let Some(a) = arr.as_any().downcast_ref::<BooleanArray>() {
123        return Value::Bool(a.value(r));
124    }
125    // Strings
126    if let Some(a) = arr.as_any().downcast_ref::<StringArray>() {
127        return Value::String(a.value(r).to_string());
128    }
129    if arr.as_any().downcast_ref::<NullArray>().is_some() {
130        return Value::Null;
131    }
132    // Fallback: stringify unknown types.
133    Value::String(format!("{arr:?}"))
134}
mongreldb_kit/arrow_util.rs

mongreldb_kit/
arrow_util.rs
