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

mongreldb_query/
arrow_conv.rs
