arrow_bridge/
lib.rs

1//! Convert [`core_query::ResultSet`] to Arrow RecordBatches and IPC stream bytes.
2
3use arrow_array::{ArrayRef, BooleanArray, Float64Array, Int64Array, RecordBatch, StringArray};
4use arrow_ipc::writer::StreamWriter;
5use arrow_schema::{DataType, Field, Schema};
6use core_query::{ResultSet, Value};
7use std::sync::Arc;
8
9#[derive(Clone, Copy, PartialEq, Eq)]
10enum ColKind {
11    Empty,
12    Int,
13    Float,
14    Bool,
15    Str,
16    Utf8,
17}
18
19fn observe(kind: ColKind, v: &Value) -> ColKind {
20    match (kind, v) {
21        (ColKind::Empty | ColKind::Int, Value::Int(_)) => ColKind::Int,
22        (ColKind::Empty | ColKind::Float, Value::Float(_)) => ColKind::Float,
23        (ColKind::Int, Value::Float(_)) | (ColKind::Float, Value::Int(_)) => ColKind::Float,
24        (ColKind::Empty | ColKind::Bool, Value::Bool(_)) => ColKind::Bool,
25        (ColKind::Empty | ColKind::Str, Value::Str(_)) => ColKind::Str,
26        (ColKind::Utf8, _) => ColKind::Utf8,
27        (_, Value::List(_)) => ColKind::Utf8,
28        _ => ColKind::Utf8,
29    }
30}
31
32fn infer_column(rs: &ResultSet, col: usize) -> ColKind {
33    let mut kind = ColKind::Empty;
34    for i in 0..rs.len() {
35        if let Some(v) = &rs.row(i)[col] {
36            kind = observe(kind, v);
37        }
38    }
39    kind
40}
41
42fn data_type(kind: ColKind) -> DataType {
43    match kind {
44        ColKind::Int => DataType::Int64,
45        ColKind::Float => DataType::Float64,
46        ColKind::Bool => DataType::Boolean,
47        ColKind::Empty | ColKind::Str | ColKind::Utf8 => DataType::Utf8,
48    }
49}
50
51fn canonical_display(v: &Value) -> String {
52    match v {
53        Value::Int(i) => i.to_string(),
54        Value::Float(f) => {
55            let s = format!("{f}");
56            if s.contains('.') || s.contains('e') || s.contains('E') {
57                s
58            } else {
59                format!("{s}.0")
60            }
61        }
62        Value::Str(s) => s.clone(),
63        Value::Bool(b) => b.to_string(),
64        Value::List(xs) => {
65            let inner: Vec<String> = xs.iter().map(canonical_display).collect();
66            format!("[{}]", inner.join(", "))
67        }
68        Value::Map(m) => {
69            let inner: Vec<String> = m
70                .iter()
71                .map(|(k, v)| format!("{k}: {}", canonical_display(v)))
72                .collect();
73            format!("{{{}}}", inner.join(", "))
74        }
75    }
76}
77
78fn as_f64(v: &Value) -> f64 {
79    match v {
80        Value::Int(i) => *i as f64,
81        Value::Float(f) => *f,
82        other => unreachable!("numeric column saw {other:?}"),
83    }
84}
85
86fn build_column(rs: &ResultSet, col: usize, kind: ColKind) -> ArrayRef {
87    match kind {
88        ColKind::Int => {
89            let vals: Vec<Option<i64>> = (0..rs.len())
90                .map(|i| {
91                    rs.row(i)[col].as_ref().map(|v| match v {
92                        Value::Int(n) => *n,
93                        other => unreachable!("int column saw {other:?}"),
94                    })
95                })
96                .collect();
97            Arc::new(Int64Array::from(vals))
98        }
99        ColKind::Float => {
100            let vals: Vec<Option<f64>> = (0..rs.len())
101                .map(|i| rs.row(i)[col].as_ref().map(as_f64))
102                .collect();
103            Arc::new(Float64Array::from(vals))
104        }
105        ColKind::Bool => {
106            let vals: Vec<Option<bool>> = (0..rs.len())
107                .map(|i| {
108                    rs.row(i)[col].as_ref().map(|v| match v {
109                        Value::Bool(b) => *b,
110                        other => unreachable!("bool column saw {other:?}"),
111                    })
112                })
113                .collect();
114            Arc::new(BooleanArray::from(vals))
115        }
116        ColKind::Empty | ColKind::Str | ColKind::Utf8 => {
117            let vals: Vec<Option<String>> = (0..rs.len())
118                .map(|i| rs.row(i)[col].as_ref().map(canonical_display))
119                .collect();
120            Arc::new(StringArray::from(vals))
121        }
122    }
123}
124
125/// Infer a per-column Arrow type and build a single-batch [`RecordBatch`].
126///
127/// Policy: all-Int → Int64; any Float among only numerics → Float64;
128/// Bool → Boolean; Str → Utf8; List or mixed types → Utf8 via canonical
129/// display; all-null → Utf8 nulls. Null cells stay null.
130pub fn to_record_batch(rs: &ResultSet) -> Result<RecordBatch, String> {
131    let kinds: Vec<ColKind> = (0..rs.columns().len())
132        .map(|c| infer_column(rs, c))
133        .collect();
134    let fields: Vec<Field> = rs
135        .columns()
136        .iter()
137        .zip(&kinds)
138        .map(|(name, kind)| Field::new(name, data_type(*kind), true))
139        .collect();
140    let schema = Arc::new(Schema::new(fields));
141    let columns: Vec<ArrayRef> = kinds
142        .into_iter()
143        .enumerate()
144        .map(|(c, kind)| build_column(rs, c, kind))
145        .collect();
146    RecordBatch::try_new(schema, columns).map_err(|e| e.to_string())
147}
148
149/// Encode `rs` as a single-batch Arrow IPC stream.
150pub fn to_ipc_bytes(rs: &ResultSet) -> Result<Vec<u8>, String> {
151    let batch = to_record_batch(rs)?;
152    let mut buf = Vec::new();
153    {
154        let mut writer =
155            StreamWriter::try_new(&mut buf, batch.schema().as_ref()).map_err(|e| e.to_string())?;
156        writer.write(&batch).map_err(|e| e.to_string())?;
157        writer.finish().map_err(|e| e.to_string())?;
158    }
159    Ok(buf)
160}
arrow_bridge/lib.rs

arrow_bridge/
lib.rs
