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    batches_to_ipc_with_checkpoint(batches, || Ok(()))
38}
39
40pub(crate) fn batches_to_ipc_controlled(
41    batches: &[RecordBatch],
42    query: &mongreldb_query::RegisteredSqlQuery,
43) -> Result<Vec<u8>> {
44    batches_to_ipc_with_checkpoint(batches, || query.checkpoint().map_err(KitError::from))
45}
46
47fn batches_to_ipc_with_checkpoint(
48    batches: &[RecordBatch],
49    mut checkpoint: impl FnMut() -> Result<()>,
50) -> Result<Vec<u8>> {
51    let schema = batches
52        .first()
53        .map(|b| b.schema())
54        .unwrap_or_else(|| std::sync::Arc::new(arrow::datatypes::Schema::empty()));
55    let mut out = Vec::new();
56    let mut writer =
57        FileWriter::try_new(&mut out, &schema).map_err(|e| KitError::Storage(e.to_string()))?;
58    for batch in batches {
59        for offset in (0..batch.num_rows()).step_by(256) {
60            checkpoint()?;
61            let length = 256.min(batch.num_rows() - offset);
62            writer
63                .write(&batch.slice(offset, length))
64                .map_err(|e| KitError::Storage(format!("arrow ipc encode: {e}")))?;
65        }
66    }
67    writer
68        .finish()
69        .map_err(|e| KitError::Storage(format!("arrow ipc encode: {e}")))?;
70    Ok(out)
71}
72
73/// Materialize one `RecordBatch` into JSON-row maps (column name → value).
74pub fn batch_to_rows(b: &RecordBatch) -> Result<Vec<Map<String, Value>>> {
75    let schema = b.schema();
76    let mut rows = Vec::with_capacity(b.num_rows());
77    for r in 0..b.num_rows() {
78        let mut row = Map::new();
79        for (c, field) in schema.fields().iter().enumerate() {
80            let name = field.name();
81            let arr = b.column(c);
82            row.insert(name.clone(), cell_value(arr.as_ref(), r));
83        }
84        rows.push(row);
85    }
86    Ok(rows)
87}
88
89/// Flatten a slice of batches into a single list of JSON-row maps.
90pub fn batches_to_rows(batches: &[RecordBatch]) -> Result<Vec<Map<String, Value>>> {
91    batches_to_rows_with_checkpoint(batches, || Ok(()))
92}
93
94pub(crate) fn batches_to_rows_controlled(
95    batches: &[RecordBatch],
96    query: &mongreldb_query::RegisteredSqlQuery,
97) -> Result<Vec<Map<String, Value>>> {
98    batches_to_rows_with_checkpoint(batches, || query.checkpoint().map_err(KitError::from))
99}
100
101fn batches_to_rows_with_checkpoint(
102    batches: &[RecordBatch],
103    mut checkpoint: impl FnMut() -> Result<()>,
104) -> Result<Vec<Map<String, Value>>> {
105    let mut rows = Vec::new();
106    for batch in batches {
107        let schema = batch.schema();
108        for row_index in 0..batch.num_rows() {
109            if row_index % 256 == 0 {
110                checkpoint()?;
111            }
112            let mut row = Map::new();
113            for (column_index, field) in schema.fields().iter().enumerate() {
114                row.insert(
115                    field.name().clone(),
116                    cell_value(batch.column(column_index).as_ref(), row_index),
117                );
118            }
119            rows.push(row);
120        }
121    }
122    Ok(rows)
123}
124
125fn cell_value(arr: &dyn Array, r: usize) -> Value {
126    if arr.is_null(r) {
127        return Value::Null;
128    }
129    // Signed integers → i64
130    if let Some(a) = arr.as_any().downcast_ref::<Int64Array>() {
131        return json!(a.value(r));
132    }
133    if let Some(a) = arr.as_any().downcast_ref::<Int32Array>() {
134        return json!(a.value(r));
135    }
136    if let Some(a) = arr.as_any().downcast_ref::<Int16Array>() {
137        return json!(a.value(r));
138    }
139    if let Some(a) = arr.as_any().downcast_ref::<Int8Array>() {
140        return json!(a.value(r));
141    }
142    // Unsigned integers → i64 (JSON has no unsigned; cast is lossless for normal values)
143    if let Some(a) = arr.as_any().downcast_ref::<UInt64Array>() {
144        return json!(a.value(r) as i64);
145    }
146    if let Some(a) = arr.as_any().downcast_ref::<UInt32Array>() {
147        return json!(a.value(r) as i64);
148    }
149    if let Some(a) = arr.as_any().downcast_ref::<UInt16Array>() {
150        return json!(a.value(r) as i64);
151    }
152    if let Some(a) = arr.as_any().downcast_ref::<UInt8Array>() {
153        return json!(a.value(r) as i64);
154    }
155    // Floats
156    if let Some(a) = arr.as_any().downcast_ref::<Float64Array>() {
157        return serde_json::Number::from_f64(a.value(r))
158            .map(Value::Number)
159            .unwrap_or(Value::Null);
160    }
161    if let Some(a) = arr.as_any().downcast_ref::<Float32Array>() {
162        return serde_json::Number::from_f64(a.value(r) as f64)
163            .map(Value::Number)
164            .unwrap_or(Value::Null);
165    }
166    // Boolean
167    if let Some(a) = arr.as_any().downcast_ref::<BooleanArray>() {
168        return Value::Bool(a.value(r));
169    }
170    // Strings
171    if let Some(a) = arr.as_any().downcast_ref::<StringArray>() {
172        return Value::String(a.value(r).to_string());
173    }
174    if arr.as_any().downcast_ref::<NullArray>().is_some() {
175        return Value::Null;
176    }
177    // Fallback: stringify unknown types.
178    Value::String(format!("{arr:?}"))
179}
mongreldb_kit/arrow_util.rs

mongreldb_kit/
arrow_util.rs
