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