mongreldb_query/
arrow_conv.rs

1//! MongrelDB ↔ Arrow conversions: schema mapping and `Vec<Row>` → `RecordBatch`.
2
3use arrow::array::{
4    ArrayRef, BooleanBuilder, FixedSizeListBuilder, Float32Builder, Float64Array, Float64Builder,
5    Int64Array, Int64Builder, StringBuilder,
6};
7use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer};
8use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
9use mongreldb_core::columnar::NativeColumn;
10use mongreldb_core::memtable::Value;
11use mongreldb_core::schema::{Schema as MongrelSchema, TypeId};
12use std::sync::Arc;
13
14use crate::error::{MongrelQueryError, Result};
15use crate::query_registry::RegisteredSqlQuery;
16
17fn bit_set(validity: &[u8], i: usize) -> bool {
18    (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1
19}
20
21/// Fast check: are all `n` positions non-null?
22fn all_bits_set(validity: &[u8], n: usize) -> bool {
23    if n == 0 {
24        return true;
25    }
26    let full = n / 8;
27    if !validity[..full].iter().all(|&b| b == 0xFF) {
28        return false;
29    }
30    if !n.is_multiple_of(8) {
31        let mask = (1u8 << (n % 8)) - 1;
32        (validity.get(full).copied().unwrap_or(0) & mask) == mask
33    } else {
34        true
35    }
36}
37
38/// Build an Arrow array straight from a typed [`NativeColumn`] (no `Value`).
39/// For the common all-non-null case on fixed-width columns, constructs the Arrow
40/// array directly from the typed buffer (one memcpy, no per-element builder).
41pub fn native_to_array(ty: TypeId, col: &NativeColumn) -> Result<ArrayRef> {
42    native_to_array_with_query(ty, col, None)
43}
44
45pub(crate) fn native_to_array_with_query(
46    ty: TypeId,
47    col: &NativeColumn,
48    query: Option<&RegisteredSqlQuery>,
49) -> Result<ArrayRef> {
50    Ok(match (ty.clone(), col) {
51        (TypeId::Int64 | TypeId::TimestampNanos, NativeColumn::Int64 { data, validity }) => {
52            if all_bits_set(validity, data.len()) {
53                Arc::new(Int64Array::new(data.clone().into(), None))
54            } else {
55                let mut b = Int64Builder::with_capacity(data.len());
56                for (i, v) in data.iter().enumerate() {
57                    checkpoint(query, i)?;
58                    if bit_set(validity, i) {
59                        b.append_value(*v);
60                    } else {
61                        b.append_null();
62                    }
63                }
64                Arc::new(b.finish())
65            }
66        }
67        (TypeId::Float64, NativeColumn::Float64 { data, validity }) => {
68            if all_bits_set(validity, data.len()) {
69                Arc::new(Float64Array::new(data.clone().into(), None))
70            } else {
71                let mut b = Float64Builder::with_capacity(data.len());
72                for (i, v) in data.iter().enumerate() {
73                    checkpoint(query, i)?;
74                    if bit_set(validity, i) {
75                        b.append_value(*v);
76                    } else {
77                        b.append_null();
78                    }
79                }
80                Arc::new(b.finish())
81            }
82        }
83        (TypeId::Bool, NativeColumn::Bool { data, validity }) => {
84            let mut b = BooleanBuilder::with_capacity(data.len());
85            for (i, v) in data.iter().enumerate() {
86                checkpoint(query, i)?;
87                if bit_set(validity, i) {
88                    b.append_value(*v != 0);
89                } else {
90                    b.append_null();
91                }
92            }
93            Arc::new(b.finish())
94        }
95        (
96            TypeId::Bytes | TypeId::Enum { .. },
97            NativeColumn::Bytes {
98                offsets,
99                values,
100                validity,
101            },
102        ) => {
103            let n = offsets.len().saturating_sub(1);
104            let mut b = StringBuilder::with_capacity(n, values.len());
105            for i in 0..n {
106                checkpoint(query, i)?;
107                if bit_set(validity, i) {
108                    let lo = offsets[i] as usize;
109                    let hi = offsets[i + 1] as usize;
110                    b.append_value(String::from_utf8_lossy(&values[lo..hi]));
111                } else {
112                    b.append_null();
113                }
114            }
115            Arc::new(b.finish())
116        }
117        _ => {
118            return Err(MongrelQueryError::Arrow(format!(
119                "native_to_array: unsupported (ty={ty:?})"
120            )))
121        }
122    })
123}
124
125/// Zero-copy variant of [`native_to_array`] for the streaming scan path. It
126/// takes ownership of the [`NativeColumn`] and, for the fixed-width `Int64` /
127/// `Float64` columns, **moves** the typed data buffer (and validity buffer when
128/// needed) straight into the Arrow array — no `memcpy`, no per-element builder.
129/// `Bool` / `Bytes` / `Embedding` fall back to the by-reference builder.
130pub fn native_to_array_owned(ty: TypeId, col: NativeColumn) -> Result<ArrayRef> {
131    native_to_array_owned_with_query(ty, col, None)
132}
133
134pub(crate) fn native_to_array_owned_with_query(
135    ty: TypeId,
136    col: NativeColumn,
137    query: Option<&RegisteredSqlQuery>,
138) -> Result<ArrayRef> {
139    query.map(RegisteredSqlQuery::checkpoint).transpose()?;
140    Ok(match (ty, col) {
141        (TypeId::Int64 | TypeId::TimestampNanos, NativeColumn::Int64 { data, validity }) => {
142            let n = data.len();
143            Arc::new(Int64Array::new(data.into(), owned_nulls(validity, n)))
144        }
145        (TypeId::Float64, NativeColumn::Float64 { data, validity }) => {
146            let n = data.len();
147            Arc::new(Float64Array::new(data.into(), owned_nulls(validity, n)))
148        }
149        // Everything else: defer to the by-reference builder.
150        (ty, col) => native_to_array_with_query(ty, &col, query)?,
151    })
152}
153
154#[inline]
155fn checkpoint(query: Option<&RegisteredSqlQuery>, index: usize) -> Result<()> {
156    if index.is_multiple_of(256) {
157        query.map(RegisteredSqlQuery::checkpoint).transpose()?;
158    }
159    Ok(())
160}
161
162/// Build an Arrow validity (`NullBuffer`) from a MongrelDB validity byte buffer,
163/// moving it without a copy. Returns `None` when every slot is non-null (Arrow
164/// treats a missing validity buffer as all-non-null). `validity` is produced by
165/// `validity_bitmap_from`, whose unused trailing bits are zero — Arrow-safe.
166fn owned_nulls(validity: Vec<u8>, n: usize) -> Option<NullBuffer> {
167    if all_bits_set(&validity, n) {
168        None
169    } else {
170        let buffer: Buffer = validity.into();
171        Some(NullBuffer::new(BooleanBuffer::new(buffer, 0, n)))
172    }
173}
174
175/// Build a `RecordBatch` directly from typed columns (vectorized scan path).
176pub fn native_columns_to_batch(
177    columns: &[(u16, NativeColumn)],
178    schema: &MongrelSchema,
179) -> Result<arrow::record_batch::RecordBatch> {
180    let mut arrays: Vec<ArrayRef> = Vec::with_capacity(schema.columns.len());
181    for cdef in &schema.columns {
182        let col = columns
183            .iter()
184            .find(|(id, _)| *id == cdef.id)
185            .map(|(_, c)| c)
186            .ok_or_else(|| MongrelQueryError::Arrow(format!("missing column {}", cdef.id)))?;
187        arrays.push(native_to_array(cdef.ty.clone(), col)?);
188    }
189    let fields: Vec<Field> = schema
190        .columns
191        .iter()
192        .map(|c| Field::new(&c.name, arrow_data_type(&c.ty).unwrap(), true))
193        .collect();
194    arrow::record_batch::RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)
195        .map_err(|e| MongrelQueryError::Arrow(e.to_string()))
196}
197
198/// Map a MongrelDB schema to an Arrow schema over the **user** columns only
199/// (system columns `_row_id`/`_epoch`/`_deleted` are hidden from SQL).
200pub fn arrow_schema(schema: &MongrelSchema) -> Result<SchemaRef> {
201    let fields: Result<Vec<Field>> = schema
202        .columns
203        .iter()
204        .map(|c| arrow_data_type(&c.ty).map(|dt| Field::new(&c.name, dt, true)))
205        .collect();
206    Ok(Arc::new(Schema::new(fields?)) as SchemaRef)
207}
208
209pub(crate) fn arrow_data_type(ty: &TypeId) -> Result<DataType> {
210    Ok(match ty {
211        TypeId::Bool => DataType::Boolean,
212        TypeId::Int8 => DataType::Int8,
213        TypeId::Int16 => DataType::Int16,
214        TypeId::Int32 | TypeId::Date32 => DataType::Int32,
215        TypeId::Int64 | TypeId::TimestampNanos => DataType::Int64,
216        TypeId::Date64 => DataType::Date64,
217        TypeId::Time64 => DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond),
218        TypeId::Interval => DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano),
219        TypeId::Uuid => DataType::FixedSizeBinary(16),
220        TypeId::Json => DataType::Utf8,
221        TypeId::Array { .. } => DataType::Utf8,
222        TypeId::UInt8 => DataType::UInt8,
223        TypeId::UInt16 => DataType::UInt16,
224        TypeId::UInt32 => DataType::UInt32,
225        TypeId::UInt64 => DataType::UInt64,
226        TypeId::Float32 => DataType::Float32,
227        TypeId::Float64 => DataType::Float64,
228        TypeId::Bytes => DataType::Utf8,
229        TypeId::Embedding { dim } => DataType::FixedSizeList(
230            Arc::new(Field::new("item", DataType::Float32, true)),
231            *dim as i32,
232        ),
233        TypeId::Decimal128 { precision, scale } => DataType::Decimal128(*precision, *scale),
234        TypeId::Enum { .. } => DataType::Utf8,
235    })
236}
237
238/// Build a single `RecordBatch` from `rows` for the user columns of `schema`.
239pub fn rows_to_batch(
240    rows: &[mongreldb_core::Row],
241    schema: &MongrelSchema,
242) -> Result<arrow::record_batch::RecordBatch> {
243    let fields: Vec<(u16, TypeId)> = schema
244        .columns
245        .iter()
246        .map(|c| (c.id, c.ty.clone()))
247        .collect();
248    let arrays: Vec<ArrayRef> = fields
249        .iter()
250        .map(|(col_id, ty)| {
251            let vals: Vec<Value> = rows
252                .iter()
253                .map(|r| r.columns.get(col_id).cloned().unwrap_or(Value::Null))
254                .collect();
255            build_array(ty.clone(), &vals)
256        })
257        .collect::<Result<_>>()?;
258    let arrow_fields: Vec<Field> = schema
259        .columns
260        .iter()
261        .map(|c| Field::new(&c.name, arrow_data_type(&c.ty).unwrap(), true))
262        .collect();
263    arrow::record_batch::RecordBatch::try_new(Arc::new(Schema::new(arrow_fields)), arrays)
264        .map_err(|e| MongrelQueryError::Arrow(e.to_string()))
265}
266
267/// Build an Arrow array from a flat slice of values (one per row).
268pub fn build_array(ty: TypeId, values: &[Value]) -> Result<ArrayRef> {
269    Ok(match ty {
270        TypeId::Int64 | TypeId::TimestampNanos => {
271            let mut b = Int64Builder::new();
272            for v in values {
273                match v {
274                    Value::Int64(x) => b.append_value(*x),
275                    _ => b.append_null(),
276                }
277            }
278            Arc::new(b.finish())
279        }
280        TypeId::Float64 => {
281            let mut b = Float64Builder::new();
282            for v in values {
283                match v {
284                    Value::Float64(x) => b.append_value(*x),
285                    _ => b.append_null(),
286                }
287            }
288            Arc::new(b.finish())
289        }
290        TypeId::Float32 => {
291            let mut b = arrow::array::Float32Builder::new();
292            for v in values {
293                match v {
294                    Value::Float64(x) => b.append_value(*x as f32),
295                    _ => b.append_null(),
296                }
297            }
298            Arc::new(b.finish())
299        }
300        TypeId::Bool => {
301            let mut b = BooleanBuilder::new();
302            for v in values {
303                match v {
304                    Value::Bool(x) => b.append_value(*x),
305                    _ => b.append_null(),
306                }
307            }
308            Arc::new(b.finish())
309        }
310        TypeId::Int32 | TypeId::Date32 => {
311            let mut b = arrow::array::Int32Builder::new();
312            for v in values {
313                match v {
314                    Value::Int64(x) => b.append_value(*x as i32),
315                    _ => b.append_null(),
316                }
317            }
318            Arc::new(b.finish())
319        }
320        TypeId::Bytes | TypeId::Enum { .. } => {
321            let mut b = StringBuilder::new();
322            for v in values {
323                match v {
324                    Value::Bytes(x) => b.append_value(String::from_utf8_lossy(x)),
325                    _ => b.append_null(),
326                }
327            }
328            Arc::new(b.finish())
329        }
330        TypeId::Embedding { dim } => {
331            let fbb = Float32Builder::new();
332            let mut b = FixedSizeListBuilder::new(fbb, dim as i32);
333            for v in values {
334                match v {
335                    Value::Embedding(x) if x.len() == dim as usize => {
336                        for fv in x {
337                            b.values().append_value(*fv);
338                        }
339                        b.append(true);
340                    }
341                    _ => {
342                        for _ in 0..dim {
343                            b.values().append_null();
344                        }
345                        b.append(false);
346                    }
347                }
348            }
349            Arc::new(b.finish())
350        }
351        TypeId::Decimal128 { precision, scale } => {
352            let mut b = arrow::array::Decimal128Builder::new()
353                .with_precision_and_scale(precision, scale)
354                .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
355            for v in values {
356                match v {
357                    Value::Decimal(d) => b.append_value(*d),
358                    _ => b.append_null(),
359                }
360            }
361            Arc::new(b.finish())
362        }
363        TypeId::Uuid => {
364            let mut b = arrow::array::FixedSizeBinaryBuilder::new(16);
365            for v in values {
366                match v {
367                    Value::Uuid(arr) => {
368                        b.append_value(arr).ok();
369                    }
370                    _ => {
371                        b.append_null();
372                    }
373                }
374            }
375            Arc::new(b.finish())
376        }
377        TypeId::Json | TypeId::Array { .. } => {
378            let mut b = arrow::array::StringBuilder::new();
379            for v in values {
380                match v {
381                    Value::Json(val) => b.append_value(String::from_utf8_lossy(val)),
382                    Value::Bytes(val) => b.append_value(String::from_utf8_lossy(val)),
383                    _ => b.append_null(),
384                }
385            }
386            Arc::new(b.finish())
387        }
388        _ => {
389            return Err(MongrelQueryError::Arrow(format!(
390                "unsupported column type {ty:?} for SQL projection"
391            )))
392        }
393    })
394}
395
396/// Build a single `RecordBatch` directly from columnar `(column_id, values)`
397/// pairs — the vectorized scan path (no row materialization).
398pub fn columns_to_batch(
399    columns: &[(u16, Vec<Value>)],
400    schema: &MongrelSchema,
401) -> Result<arrow::record_batch::RecordBatch> {
402    // Order arrays by schema column order, mapping each to its values.
403    let mut arrays: Vec<ArrayRef> = Vec::with_capacity(schema.columns.len());
404    for cdef in &schema.columns {
405        let vals = columns
406            .iter()
407            .find(|(id, _)| *id == cdef.id)
408            .map(|(_, v)| v.as_slice())
409            .unwrap_or(&[]);
410        arrays.push(build_array(cdef.ty.clone(), vals)?);
411    }
412    let arrow_fields: Vec<Field> = schema
413        .columns
414        .iter()
415        .map(|c| Field::new(&c.name, arrow_data_type(&c.ty).unwrap(), true))
416        .collect();
417    arrow::record_batch::RecordBatch::try_new(Arc::new(Schema::new(arrow_fields)), arrays)
418        .map_err(|e| MongrelQueryError::Arrow(e.to_string()))
419}
mongreldb_query/arrow_conv.rs

mongreldb_query/
arrow_conv.rs
