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

mongreldb_kit/
schema.rs
