Skip to main content

mongreldb_kit/
schema.rs

1//! Schema/value conversion between the kit model and MongrelDB core.
2
3use crate::error::{KitError, Result};
4use mongreldb_core::memtable::Value as CoreValue;
5use mongreldb_core::schema::{
6    ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema as CoreSchema, TypeId,
7};
8use mongreldb_kit_core::schema::{
9    Column, ColumnType, IndexKind as KitIndexKind, Table as KitTable,
10};
11use serde_json::{Map, Value};
12
13/// Convert a kit table to a core schema.
14pub fn to_core_schema(table: &KitTable) -> CoreSchema {
15    let columns: Vec<ColumnDef> = table
16        .columns
17        .iter()
18        .map(|c| ColumnDef {
19            id: c.id as u16,
20            name: c.name.clone(),
21            ty: match c.storage_type {
22                ColumnType::Embedding => TypeId::Embedding {
23                    dim: c.embedding_dim.unwrap_or(0),
24                },
25                other => to_core_type(other),
26            },
27            flags: to_core_flags(table, c),
28            default_value: None,
29        })
30        .collect();
31
32    let mut indexes: Vec<IndexDef> = Vec::new();
33    for idx in &table.indexes {
34        let kind = match idx.kind {
35            KitIndexKind::Bitmap => IndexKind::Bitmap,
36            KitIndexKind::Fm => IndexKind::FmIndex,
37            KitIndexKind::Ann => IndexKind::Ann,
38            KitIndexKind::Sparse => IndexKind::Sparse,
39            KitIndexKind::MinHash => IndexKind::MinHash,
40            KitIndexKind::LearnedRange => IndexKind::LearnedRange,
41        };
42        for col_name in &idx.columns {
43            if let Some(col) = table.column(col_name) {
44                indexes.push(IndexDef {
45                    name: format!("{}_{}", idx.name, col_name),
46                    column_id: col.id as u16,
47                    kind,
48                    predicate: None,
49                });
50            }
51        }
52    }
53    for uq in &table.unique_constraints {
54        for col_name in &uq.columns {
55            if let Some(col) = table.column(col_name) {
56                indexes.push(IndexDef {
57                    name: format!("uq_{}_{}", uq.name, col_name),
58                    column_id: col.id as u16,
59                    kind: IndexKind::Bitmap,
60                    predicate: None,
61                });
62            }
63        }
64    }
65
66    CoreSchema {
67        schema_id: table.id as u64,
68        columns,
69        indexes,
70        colocation: Vec::new(),
71        constraints: Default::default(),
72        clustered: false,
73    }
74}
75
76pub(crate) fn to_core_flags(table: &KitTable, column: &Column) -> ColumnFlags {
77    let mut flags = ColumnFlags::empty();
78    if column.nullable {
79        flags = flags.with(ColumnFlags::NULLABLE);
80    }
81    if table.primary_key.contains(&column.name) || column.primary_key {
82        flags = flags.with(ColumnFlags::PRIMARY_KEY);
83    }
84    if column.encrypted {
85        flags = flags.with(ColumnFlags::ENCRYPTED);
86    }
87    if column.encrypted_indexable {
88        flags = flags.with(ColumnFlags::ENCRYPTED_INDEXABLE);
89    }
90    flags
91}
92
93pub(crate) fn to_core_type(ty: ColumnType) -> TypeId {
94    match ty {
95        ColumnType::Bool => TypeId::Bool,
96        ColumnType::Int8 | ColumnType::Int16 | ColumnType::Int32 | ColumnType::Int64 => {
97            TypeId::Int64
98        }
99        ColumnType::Float32 | ColumnType::Float64 => TypeId::Float64,
100        ColumnType::Text
101        | ColumnType::Bytes
102        | ColumnType::Json
103        | ColumnType::Date
104        | ColumnType::DateTime => TypeId::Bytes,
105        ColumnType::TimestampNanos => TypeId::Int64,
106        ColumnType::Date64 => TypeId::Date64,
107        ColumnType::Time64 => TypeId::Time64,
108        ColumnType::Interval => TypeId::Interval,
109        ColumnType::Decimal128 => TypeId::Decimal128 {
110            precision: 38,
111            scale: 2,
112        },
113        ColumnType::Uuid => TypeId::Uuid,
114        ColumnType::JsonNative => TypeId::Json,
115        ColumnType::Array => TypeId::Array { element_type: 0 },
116        // Dimension is filled from the column's `embedding_dim` in
117        // `to_core_schema`; a bare type has no dimension context.
118        ColumnType::Embedding => TypeId::Embedding { dim: 0 },
119        // Sparse vectors are stored as bincode'd `Vec<(u32, f32)>` in a Bytes
120        // column; the Sparse index reads the tokens from those bytes.
121        ColumnType::Sparse => TypeId::Bytes,
122    }
123}
124
125/// Convert a JSON value to a core cell value using the column type for guidance.
126pub fn json_to_core(value: &Value, ty: ColumnType) -> Result<CoreValue> {
127    Ok(match value {
128        Value::Null => CoreValue::Null,
129        Value::Bool(b) => CoreValue::Bool(*b),
130        Value::Number(n) => {
131            if let Some(i) = n.as_i64() {
132                CoreValue::Int64(i)
133            } else {
134                CoreValue::Float64(n.as_f64().unwrap_or(f64::NAN))
135            }
136        }
137        Value::String(s) => CoreValue::Bytes(s.as_bytes().to_vec()),
138        Value::Array(arr) => {
139            if ty == ColumnType::Sparse {
140                let mut terms: Vec<(u32, f32)> = Vec::with_capacity(arr.len());
141                for pair in arr {
142                    let p = pair
143                        .as_array()
144                        .ok_or_else(|| KitError::Validation("sparse expects pairs".into()))?;
145                    let token =
146                        p.first().and_then(|v| v.as_u64()).ok_or_else(|| {
147                            KitError::Validation("sparse token must be u32".into())
148                        })? as u32;
149                    let weight = p.get(1).and_then(|v| v.as_f64()).ok_or_else(|| {
150                        KitError::Validation("sparse weight must be number".into())
151                    })? as f32;
152                    terms.push((token, weight));
153                }
154                CoreValue::Bytes(
155                    bincode::serialize(&terms).map_err(|e| KitError::Validation(e.to_string()))?,
156                )
157            } else if ty == ColumnType::Embedding {
158                let mut vec = Vec::with_capacity(arr.len());
159                for v in arr {
160                    match v.as_f64() {
161                        Some(f) => vec.push(f as f32),
162                        None => {
163                            return Err(KitError::Validation("embedding expects numbers".into()))
164                        }
165                    }
166                }
167                CoreValue::Embedding(vec)
168            } else if ty == ColumnType::Bytes {
169                let mut bytes = Vec::with_capacity(arr.len());
170                for v in arr {
171                    match v {
172                        Value::Number(n) => bytes.push(n.as_i64().unwrap_or(0) as u8),
173                        _ => return Err(KitError::Validation("bytes array expected".into())),
174                    }
175                }
176                CoreValue::Bytes(bytes)
177            } else {
178                CoreValue::Bytes(serde_json::to_vec(value)?)
179            }
180        }
181        Value::Object(_) => CoreValue::Bytes(serde_json::to_vec(value)?),
182    })
183}
184
185/// Convert a core cell value back to JSON, guided by the column type.
186pub fn core_to_json(value: &CoreValue, ty: ColumnType) -> Result<Value> {
187    Ok(match (value, ty) {
188        (CoreValue::Null, _) => Value::Null,
189        (CoreValue::Bool(b), _) => Value::Bool(*b),
190        (CoreValue::Int64(i), ColumnType::Int8) => Value::Number((*i as i8).into()),
191        (CoreValue::Int64(i), ColumnType::Int16) => Value::Number((*i as i16).into()),
192        (CoreValue::Int64(i), ColumnType::Int32) => Value::Number((*i as i32).into()),
193        (CoreValue::Int64(i), ColumnType::Int64) => Value::Number((*i).into()),
194        (CoreValue::Int64(i), ColumnType::TimestampNanos) => Value::Number((*i).into()),
195        (CoreValue::Int64(i), _) => Value::Number((*i).into()),
196        (CoreValue::Float64(f), ColumnType::Float32) => serde_json::to_value(*f as f32)?,
197        (CoreValue::Float64(f), _) => serde_json::to_value(*f)?,
198        (CoreValue::Bytes(b), ColumnType::Sparse) => {
199            let terms: Vec<(u32, f32)> =
200                bincode::deserialize(b).map_err(|e| KitError::Validation(e.to_string()))?;
201            Value::Array(
202                terms
203                    .into_iter()
204                    .map(|(t, w)| Value::Array(vec![Value::from(t), Value::from(w as f64)]))
205                    .collect(),
206            )
207        }
208        (CoreValue::Bytes(b), ColumnType::Bytes) => {
209            Value::Array(b.iter().map(|x| Value::Number((*x).into())).collect())
210        }
211        (CoreValue::Bytes(b), _) => match std::str::from_utf8(b) {
212            Ok(s) => Value::String(s.to_string()),
213            Err(_) => Value::Array(b.iter().map(|x| Value::Number((*x).into())).collect()),
214        },
215        (CoreValue::Embedding(v), _) => serde_json::to_value(v)?,
216        (CoreValue::Decimal(d), _) => Value::String(d.to_string()),
217        (
218            CoreValue::Interval {
219                months,
220                days,
221                nanos,
222            },
223            _,
224        ) => {
225            serde_json::json!({ "months": months, "days": days, "nanos": nanos })
226        }
227        (CoreValue::Uuid(b), _) => {
228            let hex: String = b.iter().map(|x| format!("{x:02x}")).collect();
229            serde_json::Value::String(hex)
230        }
231        (CoreValue::Json(b), _) => serde_json::from_slice(b.as_slice())
232            .unwrap_or_else(|_| serde_json::Value::String(String::from_utf8_lossy(b).into_owned())),
233    })
234}
235
236/// Build a JSON row from a core row and a kit table definition.
237pub fn core_row_to_json(row: &mongreldb_core::memtable::Row, table: &KitTable) -> Result<Row> {
238    let mut values = Map::new();
239    for col in &table.columns {
240        let v = row
241            .columns
242            .get(&(col.id as u16))
243            .cloned()
244            .unwrap_or(CoreValue::Null);
245        values.insert(col.name.clone(), core_to_json(&v, col.storage_type)?);
246    }
247    Ok(Row {
248        row_id: row.row_id.0,
249        values,
250    })
251}
252
253/// A kit row, identified by its internal storage row id and column values.
254#[derive(Debug, Clone, PartialEq)]
255pub struct Row {
256    pub row_id: u64,
257    pub values: Map<String, Value>,
258}
259
260impl Row {
261    /// Extract the primary-key value(s) as a JSON value.
262    ///
263    /// Single-column primary keys return the scalar value; composite keys return
264    /// an object.
265    pub fn pk(&self, table: &KitTable) -> Option<Value> {
266        if table.primary_key.len() == 1 {
267            self.values.get(&table.primary_key[0]).cloned()
268        } else {
269            let mut obj = Map::new();
270            for name in &table.primary_key {
271                obj.insert(
272                    name.clone(),
273                    self.values.get(name).cloned().unwrap_or(Value::Null),
274                );
275            }
276            Some(Value::Object(obj))
277        }
278    }
279}
280
281/// Extract the primary-key value(s) from a JSON value map.
282pub fn pk_value(values: &Map<String, Value>, table: &KitTable) -> Option<Value> {
283    if table.primary_key.len() == 1 {
284        values.get(&table.primary_key[0]).cloned()
285    } else {
286        let mut obj = Map::new();
287        for name in &table.primary_key {
288            obj.insert(
289                name.clone(),
290                values.get(name).cloned().unwrap_or(Value::Null),
291            );
292        }
293        Some(Value::Object(obj))
294    }
295}
296
297/// Convert a primary-key value into the column values for lookup.
298pub fn pk_to_map(pk: &Value, table: &KitTable) -> Result<Map<String, Value>> {
299    let mut map = Map::new();
300    match pk {
301        Value::Object(obj) => {
302            for name in &table.primary_key {
303                let v = obj
304                    .get(name)
305                    .cloned()
306                    .ok_or_else(|| KitError::Validation(format!("missing pk column {name}")))?;
307                map.insert(name.clone(), v);
308            }
309        }
310        scalar if table.primary_key.len() == 1 => {
311            map.insert(table.primary_key[0].clone(), scalar.clone());
312        }
313        _ => {
314            return Err(KitError::Validation(
315                "primary key value shape mismatch".into(),
316            ))
317        }
318    }
319    Ok(map)
320}
321
322/// Build a core cell vector from a JSON row and kit table definition.
323pub fn row_to_core_cells(
324    values: &Map<String, Value>,
325    table: &KitTable,
326) -> Result<Vec<(u16, CoreValue)>> {
327    let mut cells = Vec::with_capacity(table.columns.len());
328    for col in &table.columns {
329        let v = values.get(&col.name).cloned().unwrap_or(Value::Null);
330        cells.push((col.id as u16, json_to_core(&v, col.storage_type)?));
331    }
332    Ok(cells)
333}